yodate paper.md

Author: TomeHirata

paper.bib

@@ -1,11 +1,11 @@
 @misc{byambadalai2024estimatingdistributionaltreatmenteffects,
-      title={Estimating Distributional Treatment Effects in Randomized Experiments: Machine Learning for Variance Reduction}, 
+      title={Estimating Distributional Treatment Effects in Randomized Experiments: Machine Learning for Variance Reduction},
       author={Undral Byambadalai and Tatsushi Oka and Shota Yasui},
       year={2024},
       eprint={2407.16037},
       archivePrefix={arXiv},
       primaryClass={econ.EM},
-      url={https://arxiv.org/abs/2407.16037}, 
+      url={https://arxiv.org/abs/2407.16037},
 }
 
 @book{fisher1935design,
@@ -46,4 +46,24 @@ @article{scikit-learn
   volume={12},
   pages={2825--2830},
   year={2011}
-}
+}
+
+@misc{byambadalai2025efficientestimationdistributionaltreatment,
+      title={On Efficient Estimation of Distributional Treatment Effects under Covariate-Adaptive Randomization},
+      author={Undral Byambadalai and Tatsushi Oka and Shota Yasui},
+      year={2025},
+      eprint={2506.05945},
+      archivePrefix={arXiv},
+      primaryClass={econ.EM},
+      url={https://arxiv.org/abs/2506.05945}
+}
+
+@misc{hirata2025efficientscalableestimationdistributional,
+      title={Efficient and Scalable Estimation of Distributional Treatment Effects with Multi-Task Neural Networks},
+      author={Tomu Hirata and Undral Byambadalai and Tatsushi Oka and Shota Yasui},
+      year={2025},
+      eprint={2507.07738},
+      archivePrefix={arXiv},
+      primaryClass={econ.EM},
+      url={https://arxiv.org/abs/2507.07738}
+}

paper.md

@@ -1,71 +1,95 @@
 ---
-title: 'dte_adj: A Python package for Distributional Treatment Effects'
+title: 'dte_adj: A Python Package for Estimating Distributional Treatment Effects in Randomized Experiments'
 tags:
   - Python
-  - Distributional Treatment Effects
-  - Variance Reduction
+  - randomized experiments
+  - causal inference
+  - distributional treatment effects
+  - machine learning
+  - variance reduction
 authors:
   - name: Tomu Hirata
     orcid: 0009-0006-3140-291X
     equal-contrib: true
     affiliation: "1, 3"
-  - name: Undral Byambadalai 
+  - name: Undral Byambadalai
     corresponding: true
     affiliation: 1
   - name: Tatsushi Oka
     corresponding: true
     affiliation: "1, 2"
-  - name: Shota Yasui 
+  - name: Shota Yasui
     corresponding: true
     affiliation: 1
 affiliations:
- - name: Cyber Agent, Inc, Japan
+ - name: CyberAgent, Inc., Japan
    index: 1
  - name: Keio University, Japan
    index: 2
- - name: Indeed Technologies Japan, Japan
+ - name: Databricks Japan, Japan
    index: 3
-date: 9 August 2024
+date: 24 August 2025
 bibliography: paper.bib
-
-# Optional fields if submitting to a AAS journal too, see this blog post:
-# https://blog.joss.theoj.org/2018/12/a-new-collaboration-with-aas-publishing
-aas-doi: 10.3847/xxxxx 
-aas-journal: International Conference on Machine Learning
 ---
 
 # Summary
 
-`dte_adj` is a Python package for computing empirical cumulative distribution function (CDF) and distributional treatment effect (DTE) from data obtained by Randomized control tests. This package also contains a novel method to reduce variance of DTE using pre-treatment covariates introduced in `@Undral:2024`.
+`dte_adj` is a Python package designed for estimating distributional treatment effects (DTEs) in randomized experiments. Unlike traditional approaches that focus on average treatment effects, `dte_adj` enables researchers to analyze the full distributional impact of interventions across different outcome levels. The package implements machine learning-enhanced regression adjustment methods to achieve variance reduction, making distributional effect estimation more precise and computationally efficient. It supports multiple experimental designs including simple randomization, covariate-adaptive randomization (CAR), and local distributional treatment effect (LDTE) estimation. The package provides a scikit-learn compatible API and comprehensive functionality for computing distribution functions, probability treatment effects, and quantile treatment effects with confidence intervals.
+
+# Statement of Need
+
+Randomized experiments have been fundamental to scientific inquiry since the pioneering work of @Fisher:1935, providing the gold standard for causal inference. While most experimental analyses focus on average treatment effects (ATEs), many research questions require understanding how treatments affect the entire distribution of outcomes, not just the mean. Distributional treatment effects (DTEs) capture these richer patterns, revealing heterogeneous impacts across different outcome levels that averages can mask.
+
+Despite the growing importance of distributional analysis in fields ranging from economics to medicine, the Python ecosystem lacks comprehensive tools for DTE estimation. While SciPy provides basic empirical cumulative distribution functions, it offers no specialized functionality for treatment effect estimation, variance reduction, or confidence interval construction in experimental settings. Existing R packages like `RDDtools` focus on regression discontinuity rather than randomized experiments, and lack modern machine learning integration.
+
+`dte_adj` addresses this gap by providing a comprehensive Python framework for distributional treatment effect analysis. The package implements state-of-the-art variance reduction techniques using machine learning models for regression adjustment [@byambadalai2024estimatingdistributionaltreatmenteffects], enabling more precise DTE estimates with smaller sample sizes. It supports multiple experimental designs including covariate-adaptive randomization [@byambadalai2025efficientestimationdistributionaltreatment] and local treatment effects, with a scikit-learn [@scikit-learn] compatible API that integrates seamlessly into existing machine learning workflows. This makes advanced distributional analysis accessible to the broader Python research community, supporting more nuanced causal inference in experimental studies.
+
+# Features
+
+`dte_adj` provides a comprehensive suite of tools for distributional treatment effect analysis:
+
+## Estimator Classes
 
-# Statement of need
+The package implements multiple estimator classes following a hierarchical design pattern:
 
-Since the groundbreaking work by `@Fisher:1935`, randomized experiments have been essential in understanding the impact of interventions and shaping policy decisions. A widely used metric in this context is the Average Treatment Effect (ATE). However, exploring the distributional treatment effects often offers a more nuanced understanding than focusing solely on the average effects.
-Python is widely used in the research community recently with its flexibility and ease-of-use in the user-interface. However, there is no popular Python library for computing Distributional Treatment Effect from data obtained from randomized experiments. While scipy provides a method for computing the empirical cumulative distribution function, it lacks convenient functions for calculating DTE or for estimating the variance of the distribution.
-`dte_adj` was developed to fill the gap by offering the functionalities for 1) computing CDF from data, 2) calculating DTE and its confidence band based on CDF and 3) visualizing DTE. This library uses `numpy` as input and output of methods, which is widely used for matrix computation in Python. The main classes of this library also follows the interface of popular library `scikit-learn`, which makes it easy for the users with Machine Learning development experieneces.
+**Simple Randomization Estimators:**
+- `SimpleDistributionEstimator`: Basic empirical distribution function estimator for simple randomized experiments
+- `AdjustedDistributionEstimator`: Machine learning-enhanced estimator with regression adjustment for variance reduction
 
-# Functionalities
+**Stratified Estimators (for Covariate-Adaptive Randomization):**
+- `SimpleStratifiedDistributionEstimator`: Handles stratified block randomization designs
+- `AdjustedStratifiedDistributionEstimator`: Combines stratification with ML-based variance reduction
 
-The high level functionalities of `dte_adj` are as follows:
-1. Computing CDF and its variance based on number arrays
-2. Calculating distributional parameters and their confidence bands
-3. Visualiving distributional parameters and the confidence bands
+**Local Distribution Estimators:**
+- `SimpleLocalDistributionEstimator`: Estimates local distributional treatment effects (LDTE)
+- `AdjustedLocalDistributionEstimator`: LDTE estimation with ML adjustment for improved precision
 
-It currently offers two classes to compute CDF and its variance.
-- `SimpleDistributionEstimator`: this class offers a standard way to compute empirical CDF
-- `AdjustedDistributionEstimator`: this class offers a way to compute CDF with smaller variance adjusted by pre-treatment covariates introduced in `@Undral:2024`
+## Core Methods
 
-Both classes implement following methods to calculate distributional parameters.
-- `predict_dte`: method for computing Distributional Treatment Effect $DTE_{w, w'}(y) := F_{Y(w)}(y) - F_{Y(w')}(y)$, where $y$ is an outcome variable, $w$ is treatment type , and $F_{Y(w)}(y)$ is cumulative likelihood for treatment type $w$ and outcome $y$.
-- `predict_pte`: method for computing Probability Treatment Effect (PTE) $PTE_{w, w'}(y, h) := \left( F_{Y(w)}(y+h) - F_{Y(w)}(y) \right) - \left( F_{Y(w')}(y+h) - F_{Y(w')}(y) \right)$, where $h > 0$ is an interval of each evaluation window.
-- `predict_qte`: method for computing Quantile Treatment Effect (QTE) $QTE_{w, w'}(\tau) := F_{Y(w)}^{-1}(\tau) - F_{Y(w')}^{-1}(\tau)$, where $\tau$ is quantile.
+All estimators implement a consistent API with three primary methods:
 
-Lastly, `dte_adj.plot` module can be used for visualiting the distribution parameters. The examples of the visualization are available in the figures below.
+- `predict_dte()`: Computes Distributional Treatment Effects $DTE_{w, w'}(y) := F_{Y(w)}(y) - F_{Y(w')}(y)$, where $F_{Y(w)}(y)$ represents the cumulative distribution function for treatment $w$ at outcome level $y$.
+
+- `predict_pte()`: Computes Probability Treatment Effects over specified intervals, measuring differences in probability mass between treatment groups.
+
+- `predict_qte()`: Computes Quantile Treatment Effects $QTE_{w, w'}(\tau) := F_{Y(w)}^{-1}(\tau) - F_{Y(w')}^{-1}(\tau)$, comparing quantiles across treatments.
+
+## Advanced Features
+
+**Multi-task Learning:** The package supports multi-task neural networks (`is_multi_task=True`) for computational efficiency when analyzing many outcome locations simultaneously [@hirata2025efficientscalableestimationdistributional].
+
+**Cross-fitting:** Adjusted estimators use K-fold cross-fitting to prevent overfitting in machine learning models, ensuring robust treatment effect estimates.
+
+**Confidence Intervals:** Built-in bootstrap methods provide confidence intervals with multiple variance estimation approaches (`moment`, `simple`, `uniform`).
+
+**Visualization:** The `dte_adj.plot` module enables easy plotting of treatment effects and confidence bands.
 
 ![DTE](docs/source/_static/dte_moment.png)
-![PTE](docs/source/_static/pte_simple.png)
+![PTE](docs/source/_static/pte_empirical.png)
 ![QTE](docs/source/_static/qte.png)
 
 # Acknowledgements
 
-# References
+We thank CyberAgent, Inc. for supporting this research and the open-source community for valuable feedback during development.
+
+# References