# Attention-Based Estimation of the Individual Treatment Benefit Probability under Dose Variation

arXiv: https://arxiv.org/abs/2606.13821 · HTML: https://mlsanity.com/paper/2606.13821/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 17 of its 30 references that are in the map.
- First posted on arXiv 11 Jun 2026.

## Builds on (most cited first, 17 of 17)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Effective Approaches to Attention-based Neural Machine Translation](https://mlsanity.com/paper/1508.04025/) (2015)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms](https://mlsanity.com/paper/2101.10943/) (2021)
- [Causal Transformer for Estimating Counterfactual Outcomes](https://mlsanity.com/paper/2204.07258/) (2022)
- [Learning Counterfactual Representations for Estimating Individual Dose-Response Curves](https://mlsanity.com/paper/1902.00981/) (2019)
- [Conditional Distributional Treatment Effect with Kernel Conditional Mean Embeddings and U-Statistic Regression](https://mlsanity.com/paper/2102.08208/) (2021)
- [Conformal Prediction for Causal Effects of Continuous Treatments](https://mlsanity.com/paper/2407.03094/) (2024)
- [Quantifying Aleatoric Uncertainty of the Treatment Effect: A Novel Orthogonal Learner](https://mlsanity.com/paper/2411.03387/) (2024)
- [CETransformer: Casual Effect Estimation via Transformer Based Representation Learning](https://mlsanity.com/paper/2107.08714/) (2021)
- [Individual Treatment Effect Estimation Through Controlled Neural Network Training in Two Stages](https://mlsanity.com/paper/2201.08559/) (2022)
- [Causal Inference with Complex Treatments: A Survey](https://mlsanity.com/paper/2407.14022/) (2024)
- [Estimating Conditional Average Treatment Effects via Sufficient Representation Learning](https://mlsanity.com/paper/2408.17053/) (2024)
- [Survival Analysis as Imprecise Classification with Trainable Kernels](https://mlsanity.com/paper/2506.10140/) (2025)
- [Causal inference with multiple versions of treatment and application to personalized medicine](https://mlsanity.com/paper/2005.12427/) (2020)
- [Safe and Interpretable Estimation of Optimal Treatment Regimes](https://mlsanity.com/paper/2310.15333/) (2023)
- [Beyond the ATE: Interpretable Modelling of Treatment Effects over Dose and Time](https://mlsanity.com/paper/2507.07271/) (2025)

## Abstract (arXiv)

Estimating the probability that a treatment outperforms a control for an individual patient, called the Individual Probability of Treatment Benefit (IPTB), offers a clinically intuitive alternative to population-average metrics. However, existing methods for IPTB estimation are largely confined to binary treatment settings, despite the prevalence of dose-varying interventions in clinical practice. We propose a general framework for IPTB estimation with ordinal outcomes under discrete dose assignments, called Dose-AIPTB (Dose Attention-based IPTB). Our approach recasts the problem as binary classification over the unobserved sign of the individual treatment effect, constructing pseudo-labels from covariate-similar pairwise comparisons and aggregating them via attention mechanisms or Nadaraya-Watson kernel regression. This formulation naturally accommodates multiple discrete dose levels, extending beyond the binary treatment paradigm. Through numerical experiments on real-world and synthetic data under covariate shift, varying sample sizes, and heterogeneous outcomes, we demonstrate that attention-based aggregation consistently outperforms kernel alternatives. The framework provides a foundation for personalized dose selection grounded in individual-level benefit probabilities. Codes implementing the model are publicly available at https://github.com/NTAILab/AIPTBDose.