Fixed-Confidence Best-Arm Identification for Causal Mediation Analysis

Harsh Shrivastava, Yuta Kawakami, Junpei Komiyama, Jin Tian
Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, PMLR 337:6283-6309, 2026.

Abstract

This paper studies the problem of identifying the treatment that maximizes the expected natural direct potential outcome (NDPO), which captures the potential outcome of an intervention while excluding the pathway transmitted through a mediator that researchers may wish to remove from evaluation. We first establish population-level identification of the expected NDPO in a causal bandit setting using observable interventional distributions. We then develop a fixed-confidence best-arm identification (BAI) algorithm based on the Track-and-Stop (TaS) framework, employing a cutting-set method to solve the resulting semi-infinite optimization problem. The proposed algorithm achieves sample-efficient identification with a high-probability correctness guarantee. We prove that it satisfies $\delta$-correctness and asymptotic optimality. Finally, we validate the approach through empirical evaluations on a large-scale real-world advertising dataset (IPinYou).

Cite this Paper

BibTeX
@InProceedings{pmlr-v337-shrivastava26a, title = {Fixed-Confidence Best-Arm Identification for Causal Mediation Analysis}, author = {Shrivastava, Harsh and Kawakami, Yuta and Komiyama, Junpei and Tian, Jin}, booktitle = {Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence}, pages = {6283--6309}, year = {2026}, editor = {Perković, Emilija and Malinsky, Daniel}, volume = {337}, series = {Proceedings of Machine Learning Research}, month = {17--21 Aug}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v337/main/assets/shrivastava26a/shrivastava26a.pdf}, url = {https://proceedings.mlr.press/v337/shrivastava26a.html}, abstract = {This paper studies the problem of identifying the treatment that maximizes the expected natural direct potential outcome (NDPO), which captures the potential outcome of an intervention while excluding the pathway transmitted through a mediator that researchers may wish to remove from evaluation. We first establish population-level identification of the expected NDPO in a causal bandit setting using observable interventional distributions. We then develop a fixed-confidence best-arm identification (BAI) algorithm based on the Track-and-Stop (TaS) framework, employing a cutting-set method to solve the resulting semi-infinite optimization problem. The proposed algorithm achieves sample-efficient identification with a high-probability correctness guarantee. We prove that it satisfies $\delta$-correctness and asymptotic optimality. Finally, we validate the approach through empirical evaluations on a large-scale real-world advertising dataset (IPinYou).} }
Endnote
%0 Conference Paper %T Fixed-Confidence Best-Arm Identification for Causal Mediation Analysis %A Harsh Shrivastava %A Yuta Kawakami %A Junpei Komiyama %A Jin Tian %B Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence %C Proceedings of Machine Learning Research %D 2026 %E Emilija Perković %E Daniel Malinsky %F pmlr-v337-shrivastava26a %I PMLR %P 6283--6309 %U https://proceedings.mlr.press/v337/shrivastava26a.html %V 337 %X This paper studies the problem of identifying the treatment that maximizes the expected natural direct potential outcome (NDPO), which captures the potential outcome of an intervention while excluding the pathway transmitted through a mediator that researchers may wish to remove from evaluation. We first establish population-level identification of the expected NDPO in a causal bandit setting using observable interventional distributions. We then develop a fixed-confidence best-arm identification (BAI) algorithm based on the Track-and-Stop (TaS) framework, employing a cutting-set method to solve the resulting semi-infinite optimization problem. The proposed algorithm achieves sample-efficient identification with a high-probability correctness guarantee. We prove that it satisfies $\delta$-correctness and asymptotic optimality. Finally, we validate the approach through empirical evaluations on a large-scale real-world advertising dataset (IPinYou).
APA
Shrivastava, H., Kawakami, Y., Komiyama, J. & Tian, J.. (2026). Fixed-Confidence Best-Arm Identification for Causal Mediation Analysis. Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, in Proceedings of Machine Learning Research 337:6283-6309 Available from https://proceedings.mlr.press/v337/shrivastava26a.html.

Related Material