Technical Summary: Fixed-Confidence Best-Arm Identification for Causal Mediation Analysis

1. Problem Definition

This paper addresses the problem of Best-Arm Identification (BAI) within a causal bandit setting, specifically focusing on causal mediation analysis.

• Context: In standard causal bandits, the objective is typically to maximize the interventional mean $E[Y | do(X=x)]$, which aggregates effects across all causal pathways. However, researchers often wish to isolate the effect of a treatment $X$ on an outcome $Y$ that does not pass through a specific mediator $Z$ (e.g., to exclude unethical mechanisms, positioning bias in ads, or adverse intermediate reactions).
• Objective: The paper aims to identify the arm $x^*$ that maximizes the Expected Natural Direct Potential Outcome (NDPO), defined as $E[Y_{x, Z_{x_0}}]$. This quantity represents the expected outcome if the treatment is set to $x$ while the mediator $Z$ is held at the value it would have taken under a reference arm $x_0$.
• Setting: The problem is formulated in a fixed-confidence setting. A learner sequentially pulls arms (interventions $do(X=x)$), observing i.i.d. samples from the interventional distribution $P(Z, Y | do(X=x))$. The goal is to identify the optimal arm with probability at least $1-\delta$ while minimizing the expected stopping time (sample complexity).

2. Methodology

2.1 Population-Level Identification

The authors first establish that the expected NDPO is identifiable from observable interventional distributions under a causal bandit setting.

• Assumption: They rely on the standard identification assumption for natural direct effects: $Y_{x,z} \perp\!\!\perp Z_{x'}$ for all $x, x', z$. Graphically, this corresponds to the absence of unmeasured confounding between the mediator $Z$ and outcome $Y$ in the graph where outgoing edges from $X$ and $Z$ are removed.
• Identification Formula: Under this assumption, the expected NDPO $\theta(x)$ is identified as:
  $$\theta(x) = \sum_{y,z} y P(Y=y | Z=z, do(X=x)) P(Z=z | do(X=x_0))$$
  This formulation allows estimation using only the joint interventional distributions $P(Z, Y | do(X=x))$ and $P(Z | do(X=x_0))$, avoiding the need for nested counterfactuals that are not directly observable in the bandit setting.

2.2 Sample Complexity Lower Bound

The paper derives an instance-dependent lower bound on the sample complexity for any $\delta$-correct algorithm.

• Key Difference from Standard BAI: Unlike standard BAI where the objective depends only on the arm's own distribution, the NDPO couples the outcome distribution under arm $x$ with the mediator distribution under the baseline arm $x_0$.
• Complexity Structure: The lower bound involves a semi-infinite optimization problem. The characteristic time $T^*(P)$ is defined by:
  $$(T^*(P))^{-1} = \sup_{w \in \Sigma_K} \inf_{Q \in \text{Alt}(P)} \sum_{x \in \mathcal{X}} w_x KL(P_x \| Q_x)$$
  where $\text{Alt}(P)$ is the set of alternative models where a different arm is optimal. The constraint defining $\text{Alt}(P)$ is non-separable because the baseline mediator distribution $Q_{x_0}(Z)$ appears in the objective for every arm, creating a coupling across arms that does not exist in classical BAI.

2.3 Proposed Algorithm: TaS-NDPO

The authors propose TaS-NDPO, a Track-and-Stop (TaS) algorithm adapted for this specific objective. The algorithm consists of three main components:

1. Plug-in Estimation: Empirical estimates of the interventional distributions are used to compute plug-in estimates of $\theta(x)$.
2. Sampling Rule (D-tracking with Forcing):
   • D-tracking: The algorithm tracks the optimal sampling proportions $w^*$ derived from the lower bound.
   • Competitor Coverage: To handle the non-separable constraints, the algorithm ensures every arm is selected as a "competitor" in the optimization problem sufficiently often.
   • Cell-level Forcing: Crucially, the algorithm forces exploration at the cell level $(x, z)$ (combinations of arm and mediator value). This is necessary because the KL divergence decomposes into terms weighted by mediator probabilities; rare mediator values must be sampled enough to ensure consistent estimation of conditional outcome probabilities $P(Y|Z, do(X))$.
3. Stopping Rule (GLRT): The algorithm stops when the Generalized Likelihood Ratio Test (GLRT) statistic exceeds a threshold $\beta(t, \delta)$. The statistic measures the minimum KL divergence to any alternative model where a suboptimal arm appears better than the current empirical best.
4. Optimization Solver: Solving the inner minimization in the sampling rule and stopping rule involves a semi-infinite optimization problem. The authors employ a cutting-set method (Mutapcic and Boyd, 2009) combined with a reduction to three arms (the baseline, the current best, and a competitor) and an alternating minimization scheme to solve the resulting bi-convex subproblems.

3. Key Contributions

1. Problem Formulation: The paper formally defines and studies the BAI problem for maximizing the Expected NDPO, a quantity relevant for isolating direct causal effects in the presence of mediators.
2. Identification in Bandits: It establishes the identification of the expected NDPO using only interventional distributions available in a causal bandit setting, bridging causal inference theory with sequential decision-making.
3. Theoretical Bounds: It derives a new sample complexity lower bound that highlights the unique challenges of NDPO optimization, specifically the coupling of arms via the baseline mediator distribution and the need for cell-level information accumulation.
4. Algorithm Design: It introduces TaS-NDPO, a novel algorithm that:
   • Uses a cutting-set method to solve the complex semi-infinite optimization required for the optimal sampling proportions.
   • Incorporates cell-level forcing to ensure consistency in sparse mediator regimes.
   • Proves $\delta$-correctness (guaranteed correctness with probability $1-\delta$) and asymptotic optimality (achieving the lower bound on sample complexity as $\delta \to 0$).

4. Experimental Results

The authors evaluated TaS-NDPO on both synthetic data and a large-scale real-world advertising dataset (IPinYou).

• Real-World Validation (IPinYou):
  • The dataset involved selecting ad creatives ($X$) to maximize clicks ($Y$), with slot visibility ($Z$) as the mediator.
  • Decision Divergence: The arm maximizing the standard interventional mean was different from the arm maximizing the NDPO. The standard mean favored a creative that appeared frequently in top slots (mediated effect), while the NDPO-optimal arm was different.
  • Performance: TaS-NDPO achieved a median stopping time of 12,928 samples, roughly 50% more efficient than the arm-level variant (25,856 samples) and significantly outperforming uniform sampling and a baseline-first approach. All methods maintained zero empirical error rates at $\delta=0.05$.
• Synthetic Experiments:
  • Gap Scaling: Results confirmed the $1/\Delta^2$ scaling of sample complexity predicted by the lower bound.
  • Mediator Sparsity: The advantage of cell-level forcing was most pronounced in settings with rare mediator-outcome cells (sparse regimes), where arm-level forcing failed to accumulate sufficient information.
  • Runtime: The computational overhead of the cutting-set optimization was moderate (average 22.5 seconds per run for up to 100k samples), not constituting a bottleneck compared to the gains in sample efficiency.

5. Significance and Claims

The paper claims that TaS-NDPO provides a theoretically grounded and empirically efficient solution for identifying treatments that maximize direct causal effects while excluding mediated pathways.

• Theoretical Significance: It extends the Track-and-Stop framework to causal settings with nested counterfactuals, proving that asymptotic optimality is achievable even when the objective function couples distributions across different arms (via the baseline mediator).
• Practical Significance: By isolating direct effects, the method enables decision-making that avoids undesirable mechanisms (e.g., clickbait, adverse side effects, or unethical positioning) which might otherwise be favored by standard mean-maximization approaches.
• Modesty: The authors acknowledge that their current framework focuses on binary outcomes and a specific causal graph structure (Mediator between Treatment and Outcome). They note that extending to multi-category outcomes is straightforward, but generalizing to more complex path-specific effects in broader causal structures remains a direction for future work. They also clarify that the identification relies on the plausibility of the underlying causal graph (specifically the absence of unmeasured confounding between $Z$ and $Y$), which is a standard but untestable assumption in causal inference.