IDEAS home Printed from https://ideas.repec.org/p/arx/papers/2507.07738.html
   My bibliography  Save this paper

Efficient and Scalable Estimation of Distributional Treatment Effects with Multi-Task Neural Networks

Author

Listed:
  • Tomu Hirata
  • Undral Byambadalai
  • Tatsushi Oka
  • Shota Yasui
  • Shingo Uto

Abstract

We propose a novel multi-task neural network approach for estimating distributional treatment effects (DTE) in randomized experiments. While DTE provides more granular insights into the experiment outcomes over conventional methods focusing on the Average Treatment Effect (ATE), estimating it with regression adjustment methods presents significant challenges. Specifically, precision in the distribution tails suffers due to data imbalance, and computational inefficiencies arise from the need to solve numerous regression problems, particularly in large-scale datasets commonly encountered in industry. To address these limitations, our method leverages multi-task neural networks to estimate conditional outcome distributions while incorporating monotonic shape constraints and multi-threshold label learning to enhance accuracy. To demonstrate the practical effectiveness of our proposed method, we apply our method to both simulated and real-world datasets, including a randomized field experiment aimed at reducing water consumption in the US and a large-scale A/B test from a leading streaming platform in Japan. The experimental results consistently demonstrate superior performance across various datasets, establishing our method as a robust and practical solution for modern causal inference applications requiring a detailed understanding of treatment effect heterogeneity.

Suggested Citation

  • Tomu Hirata & Undral Byambadalai & Tatsushi Oka & Shota Yasui & Shingo Uto, 2025. "Efficient and Scalable Estimation of Distributional Treatment Effects with Multi-Task Neural Networks," Papers 2507.07738, arXiv.org.
  • Handle: RePEc:arx:papers:2507.07738
    as

    Download full text from publisher

    File URL: http://arxiv.org/pdf/2507.07738
    File Function: Latest version
    Download Restriction: no
    ---><---

    References listed on IDEAS

    as
    1. Victor Chernozhukov & Iván Fernández-Val & Blaise Melly & Kaspar Wüthrich, 2020. "Generic Inference on Quantile and Quantile Effect Functions for Discrete Outcomes," Journal of the American Statistical Association, Taylor & Francis Journals, vol. 115(529), pages 123-137, January.
    2. Sergio Firpo, 2007. "Efficient Semiparametric Estimation of Quantile Treatment Effects," Econometrica, Econometric Society, vol. 75(1), pages 259-276, January.
    3. Markus Frölich & Blaise Melly, 2013. "Unconditional Quantile Treatment Effects Under Endogeneity," Journal of Business & Economic Statistics, Taylor & Francis Journals, vol. 31(3), pages 346-357, July.
    4. Victor Chernozhukov & Juan Carlos Escanciano & Hidehiko Ichimura & Whitney K. Newey & James M. Robins, 2022. "Locally Robust Semiparametric Estimation," Econometrica, Econometric Society, vol. 90(4), pages 1501-1535, July.
    5. Undral Byambadalai & Tomu Hirata & Tatsushi Oka & Shota Yasui, 2025. "On Efficient Estimation of Distributional Treatment Effects under Covariate-Adaptive Randomization," Papers 2506.05945, arXiv.org.
    6. Victor Chernozhukov & Christian Hansen, 2005. "An IV Model of Quantile Treatment Effects," Econometrica, Econometric Society, vol. 73(1), pages 245-261, January.
    7. Fortin, Nicole & Lemieux, Thomas & Firpo, Sergio, 2011. "Decomposition Methods in Economics," Handbook of Labor Economics, in: O. Ashenfelter & D. Card (ed.), Handbook of Labor Economics, edition 1, volume 4, chapter 1, pages 1-102, Elsevier.
    8. A. Belloni & V. Chernozhukov & I. Fernández‐Val & C. Hansen, 2017. "Program Evaluation and Causal Inference With High‐Dimensional Data," Econometrica, Econometric Society, vol. 85, pages 233-298, January.
    9. Undral Byambadalai & Tatsushi Oka & Shota Yasui, 2024. "Estimating Distributional Treatment Effects in Randomized Experiments: Machine Learning for Variance Reduction," Papers 2407.16037, arXiv.org.
    10. Susan Athey & Guido W. Imbens, 2006. "Identification and Inference in Nonlinear Difference-in-Differences Models," Econometrica, Econometric Society, vol. 74(2), pages 431-497, March.
    11. repec:cdl:ucsdec:qt5zm6m9rq is not listed on IDEAS
    12. John A. List & Ian Muir & Gregory Sun, 2024. "Using machine learning for efficient flexible regression adjustment in economic experiments," Econometric Reviews, Taylor & Francis Journals, vol. 44(1), pages 2-40, July.
    13. Brantly Callaway & Tong Li, 2019. "Quantile treatment effects in difference in differences models with panel data," Quantitative Economics, Econometric Society, vol. 10(4), pages 1579-1618, November.
    14. Abadie A., 2002. "Bootstrap Tests for Distributional Treatment Effects in Instrumental Variable Models," Journal of the American Statistical Association, American Statistical Association, vol. 97, pages 284-292, March.
    15. Rothe, Christoph, 2010. "Nonparametric estimation of distributional policy effects," Journal of Econometrics, Elsevier, vol. 155(1), pages 56-70, March.
    16. Dandan Xu & Michael J. Daniels & Almut G. Winterstein, 2018. "A Bayesian nonparametric approach to causal inference on quantiles," Biometrics, The International Biometric Society, vol. 74(3), pages 986-996, September.
    17. Callaway, Brantly & Li, Tong & Oka, Tatsushi, 2018. "Quantile treatment effects in difference in differences models under dependence restrictions and with only two time periods," Journal of Econometrics, Elsevier, vol. 206(2), pages 395-413.
    Full references (including those not matched with items on IDEAS)

    Most related items

    These are the items that most often cite the same works as this one and are cited by the same works as this one.
    1. Undral Byambadalai & Tatsushi Oka & Shota Yasui, 2024. "Estimating Distributional Treatment Effects in Randomized Experiments: Machine Learning for Variance Reduction," Papers 2407.16037, arXiv.org.
    2. Pedro H. C. Sant'Anna, 2016. "Program Evaluation with Right-Censored Data," Papers 1604.02642, arXiv.org.
    3. Brantly Callaway & Weige Huang, 2020. "Distributional Effects of a Continuous Treatment with an Application on Intergenerational Mobility," Oxford Bulletin of Economics and Statistics, Department of Economics, University of Oxford, vol. 82(4), pages 808-842, August.
    4. Undral Byambadalai & Tomu Hirata & Tatsushi Oka & Shota Yasui, 2025. "On Efficient Estimation of Distributional Treatment Effects under Covariate-Adaptive Randomization," Papers 2506.05945, arXiv.org.
    5. Tatsushi Oka & Shota Yasui & Yuta Hayakawa & Undral Byambadalai, 2024. "Regression Adjustment for Estimating Distributional Treatment Effects in Randomized Controlled Trials," Papers 2407.14074, arXiv.org, revised Jan 2025.
    6. Pedro H. C. Sant'Anna & Xiaojun Song & Qi Xu, 2022. "Covariate distribution balance via propensity scores," Journal of Applied Econometrics, John Wiley & Sons, Ltd., vol. 37(6), pages 1093-1120, September.
    7. Victor Chernozhukov & Iván Fernández‐Val & Blaise Melly, 2013. "Inference on Counterfactual Distributions," Econometrica, Econometric Society, vol. 81(6), pages 2205-2268, November.
    8. Callaway, Brantly & Li, Tong & Oka, Tatsushi, 2018. "Quantile treatment effects in difference in differences models under dependence restrictions and with only two time periods," Journal of Econometrics, Elsevier, vol. 206(2), pages 395-413.
    9. Afrouz Azadikhah Jahromi & Brantly Callaway, 2022. "Heterogeneous Effects of Job Displacement on Earnings," Empirical Economics, Springer, vol. 62(1), pages 213-245, January.
    10. Zequn Jin & Lihua Lin & Zhengyu Zhang, 2022. "Identification and Auto-debiased Machine Learning for Outcome Conditioned Average Structural Derivatives," Papers 2211.07903, arXiv.org.
    11. Kaspar Wüthrich, 2020. "A Comparison of Two Quantile Models With Endogeneity," Journal of Business & Economic Statistics, Taylor & Francis Journals, vol. 38(2), pages 443-456, April.
    12. repec:cdl:ucsdec:qt0q43931f is not listed on IDEAS
    13. Volpe Martincus, Christian & Carballo, Jerónimo, 2010. "Beyond the average effects: The distributional impacts of export promotion programs in developing countries," Journal of Development Economics, Elsevier, vol. 92(2), pages 201-214, July.
    14. Guido W. Imbens & Jeffrey M. Wooldridge, 2009. "Recent Developments in the Econometrics of Program Evaluation," Journal of Economic Literature, American Economic Association, vol. 47(1), pages 5-86, March.
    15. Neng-Chieh Chang, 2020. "The Mode Treatment Effect," Papers 2007.11606, arXiv.org.
    16. Edouard Pignède, 2025. "Who carries the burden of climate change? Heterogeneous impact of droughts in sub‐Saharan Africa," American Journal of Agricultural Economics, John Wiley & Sons, vol. 107(3), pages 925-957, May.
    17. Yang Ning & Sida Peng & Jing Tao, 2020. "Doubly Robust Semiparametric Difference-in-Differences Estimators with High-Dimensional Data," Papers 2009.03151, arXiv.org.
    18. Damian Clarke & Manuel Llorca Jaña & Daniel Pailañir, 2023. "The use of quantile methods in economic history," Historical Methods: A Journal of Quantitative and Interdisciplinary History, Taylor & Francis Journals, vol. 56(2), pages 115-132, April.
    19. Christoph Rothe, 2012. "Partial Distributional Policy Effects," Econometrica, Econometric Society, vol. 80(5), pages 2269-2301, September.
    20. Callaway, Brantly & Sant’Anna, Pedro H.C., 2021. "Difference-in-Differences with multiple time periods," Journal of Econometrics, Elsevier, vol. 225(2), pages 200-230.
    21. Takuya Ishihara, 2020. "Panel Data Quantile Regression for Treatment Effect Models," Papers 2001.04324, arXiv.org, revised Nov 2021.

    More about this item

    NEP fields

    This paper has been announced in the following NEP Reports:

    Statistics

    Access and download statistics

    Corrections

    All material on this site has been provided by the respective publishers and authors. You can help correct errors and omissions. When requesting a correction, please mention this item's handle: RePEc:arx:papers:2507.07738. See general information about how to correct material in RePEc.

    If you have authored this item and are not yet registered with RePEc, we encourage you to do it here. This allows to link your profile to this item. It also allows you to accept potential citations to this item that we are uncertain about.

    If CitEc recognized a bibliographic reference but did not link an item in RePEc to it, you can help with this form .

    If you know of missing items citing this one, you can help us creating those links by adding the relevant references in the same way as above, for each refering item. If you are a registered author of this item, you may also want to check the "citations" tab in your RePEc Author Service profile, as there may be some citations waiting for confirmation.

    For technical questions regarding this item, or to correct its authors, title, abstract, bibliographic or download information, contact: arXiv administrators (email available below). General contact details of provider: http://arxiv.org/ .

    Please note that corrections may take a couple of weeks to filter through the various RePEc services.

    IDEAS is a RePEc service. RePEc uses bibliographic data supplied by the respective publishers.