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

2026-06-11

AI Summary

This study addresses the limitation of existing methods that are largely confined to binary treatment settings and struggle to accurately estimate individual probability of treatment benefit (IPTB) under discrete multi-dose interventions. The authors reformulate IPTB estimation as a binary classification problem targeting the sign of individual treatment effects. They construct pseudo-labels using pairs of covariate-similar samples and, for the first time, introduce an attention mechanism to aggregate information, replacing conventional Nadaraya-Watson kernel regression. This approach overcomes the binary treatment constraint and enables flexible modeling of personalized benefit probabilities across multiple doses. Experimental results on both real-world and synthetic data demonstrate that the proposed attention-based aggregation consistently outperforms kernel methods under various challenging conditions—including covariate shift, varying sample sizes, and heterogeneous treatment responses—providing a robust foundation for personalized dose decision-making.

Abstract

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.

Problem

• Individual Treatment Benefit
• Dose Variation
• IPTB Estimation
• Ordinal Outcomes
• Personalized Dosing

Innovation

• Individual Treatment Benefit
• Attention Mechanism
• Dose Response
• Pseudo-labeling
• Personalized Medicine