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

# Attention-Based Estimation of the Individual Treatment Benefit Probability under Dose Variation
Source: [https://arxiv.org/html/2606.13821](https://arxiv.org/html/2606.13821)
11institutetext:Peter the Great St\.Petersburg Polytechnic University
Higher School of Artificial Intelligence Technologies
St\.Petersburg, Russia
11email:utkin\_lv@spbstu\.ru, konstantinov\_av@spbstu\.ru, kogan\_sk@spbstu\.ru, verbova\_nm@spbstu\.ru, goryunov\_myu@spbstu\.ruAndrei V\. KonstantinovStanislav K\. KoganNatalya M\. VerbovaMaksim I\. Goriunov

###### 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](https://github.com/NTAILab/AIPTBDose)\.

## 1Introduction

Clinical decision\-making has long been guided by evidence generated at the population level\. The*Average Treatment Effect*\(ATE\) stands as the cornerstone of this paradigm, offering a singular measure that determines whether a therapy outperforms a control for a prototypical patient\. This metric has proven indispensable for regulatory approval and the formulation of clinical guidelines\. Yet, a fundamental tension arises when the ATE is transported from the research setting to the bedside\. Whereas the ATE resolves uncertainty about group\-level superiority, the clinician faces a different question:*Will this specific patient, with this particular constellation of characteristics, derive benefit?*The principle of clinical equipoise, which justifies randomization at the trial level, offers little guidance for the individual case\.

This disconnect has spurred a shift toward individualization\. The Conditional Average Treatment Effect \(CATE\) emerged as a natural extension, aiming to characterize how treatment effects vary across subgroups defined by patient covariates\. However, even CATE, which estimates the*expected*difference in outcomes, falls short of addressing the clinician’s core concern\. What matters in practice is not merely the expected magnitude of benefit, but the probability that a patient will experience a favorable outcome relative to the alternative\. This motivates the*Individual Probability of Treatment Benefit*\(IPTB\), defined for a patient with covariatesX=𝐱X=\\mathbf\{x\}as

ρ​\(𝐱\)=Pr⁡\{H\>Y∣X=𝐱\}=Pr⁡\{Δ\>0∣X=𝐱\},\\rho\(\\mathbf\{x\}\)=\\Pr\\\{H\>Y\\mid X=\\mathbf\{x\}\\\}=\\Pr\\\{\\Delta\>0\\mid X=\\mathbf\{x\}\\\},\(1\)whereHHandYYdenote the potential outcomes under treatment and control, respectively, andΔ=H−Y\\Delta=H\-Ycaptures the individual\-level treatment effect\. Unlike CATE, which summarizes the distribution ofΔ\\Deltathrough its first moment, the IPTB directly targets the probability of a positive outcome as a quantity that aligns directly with the notion of treatment benefit in clinical reasoning\.

The appeal of the IPTB belies its statistical complexity\. Estimating a mean effect \(whether unconditional \(ATE\) or conditional \(CATE\)\) requires only the ability to model the conditional expectation of the observed outcome\. By contrast, estimatingPr⁡\{Δ\>0∣X=𝐱\}\\Pr\\\{\\Delta\>0\\mid X=\\mathbf\{x\}\\\}demands recovering the full conditional distribution of potential outcomes, or at least the joint distribution of\(H,Y\)\(H,Y\)given covariates\. This task is fundamentally complicated by the central obstacle of causal inference: for any given individual, at most one potential outcome is ever observed\. The counterfactual outcome remains hidden, making the sign ofΔ\\Deltainherently unobservable at the individual level\.

Reconstructing the distribution ofΔ\\Deltathus requires either parametric assumptions about the underlying data\-generating process or sophisticated nonparametric strategies capable of imputing the missing counterfactual distribution\. Bayesian approaches naturally accommodate this challenge by modeling joint potential outcome distributions\[[31](https://arxiv.org/html/2606.13821#bib.bib31)\], while frequentist methods have increasingly turned to techniques such as probabilistic classification or direct modeling of the benefit function\[[14](https://arxiv.org/html/2606.13821#bib.bib14),[13](https://arxiv.org/html/2606.13821#bib.bib13)\]\. What unifies these efforts is the recognition that inferring individual benefit probabilities imposes stronger demands both in terms of assumptions and computational complexity than conventional effect estimation\.

The literature on heterogeneous treatment effect estimation has witnessed substantial methodological evolution\. Early approaches relied on regularized regression with interaction terms\[[8](https://arxiv.org/html/2606.13821#bib.bib8)\]\. Subsequent advances introduced meta\-learners, including the T\-learner, S\-learner, X\-learner, and DR\-learner\[[9](https://arxiv.org/html/2606.13821#bib.bib9)\], that combine base estimators in flexible ways to estimate CATE\[[11](https://arxiv.org/html/2606.13821#bib.bib11),[21](https://arxiv.org/html/2606.13821#bib.bib21),[27](https://arxiv.org/html/2606.13821#bib.bib27)\]\. Deep learning architectures have since been deployed to capture complex, high\-dimensional relationships in covariate spaces\[[3](https://arxiv.org/html/2606.13821#bib.bib3),[16](https://arxiv.org/html/2606.13821#bib.bib16),[20](https://arxiv.org/html/2606.13821#bib.bib20),[24](https://arxiv.org/html/2606.13821#bib.bib24)\]\. Nonparametric kernel methods\[[7](https://arxiv.org/html/2606.13821#bib.bib7),[18](https://arxiv.org/html/2606.13821#bib.bib18)\]and transformer\-based attention mechanisms\[[4](https://arxiv.org/html/2606.13821#bib.bib4),[14](https://arxiv.org/html/2606.13821#bib.bib14),[30](https://arxiv.org/html/2606.13821#bib.bib30)\]represent more recent frontiers, offering flexible function approximation with varying trade\-offs in interpretability and scalability\.

Parallel developments have targeted IPTB estimation specifically\. These include Bayesian frameworks that model the joint distribution of potential outcomes\[[31](https://arxiv.org/html/2606.13821#bib.bib31)\], as well as approaches that reframe the problem as probabilistic classification, where the target is the probability that the treated outcome exceeds the control outcome\[[14](https://arxiv.org/html/2606.13821#bib.bib14),[13](https://arxiv.org/html/2606.13821#bib.bib13)\]\. Despite these advances, a common limitation persists across much of this literature: the implicit assumption of binary treatment assignment\.

In many clinical contexts, treatment is not simply present or absent; it is administered in varying intensities, frequencies, or quantities\. Whether conceptualized as the number of drug units, the duration of therapy, or the concentration of an active ingredient, the*dose*introduces a layer of complexity that binary formulations cannot accommodate\. A patient may derive benefit from a low dose, experience toxicity at a high dose, or exhibit a non\-monotonic response that defies simple dose response assumptions\. Moreover, the optimal dose defined as the minimum level sufficient to achieve a positive treatment effect may vary across individuals in ways that reflect underlying biological heterogeneity\.

This gap has not gone unnoticed\. Recent work has begun exploring causal inference with complex treatments, including continuous and multi\-valued interventions\[[6](https://arxiv.org/html/2606.13821#bib.bib6),[17](https://arxiv.org/html/2606.13821#bib.bib17),[19](https://arxiv.org/html/2606.13821#bib.bib19),[22](https://arxiv.org/html/2606.13821#bib.bib22),[28](https://arxiv.org/html/2606.13821#bib.bib28)\]\. The problem of multiple treatment versions has been examined from a causal identification perspective\[[1](https://arxiv.org/html/2606.13821#bib.bib1),[25](https://arxiv.org/html/2606.13821#bib.bib25)\]\. Notably, neural network\-based approaches have been proposed for estimating individual dose response curves across continuous dosage parameters\[[23](https://arxiv.org/html/2606.13821#bib.bib23)\]\. However, these methods typically target expected outcomes or dose response surfaces rather than the probability of benefit a quantity that may better capture clinically meaningful thresholds of efficacy\.

In this work, we introduce a general framework for estimating the IPTB with continuous or ordinal outcomes in settings where treatment is administered at discrete dose levels\. It is called Dose\-AIPTB \(Dose Attention\-based IPTB\)\. Our approach recasts the estimation problem as binary classification, where the target label is the sign of the individual treatment effectΔ\\Delta\. SinceΔ\\Deltais unobservable, we construct pseudo\-labels from pairwise comparisons between treated and control patients with similar covariate profiles\. These comparisons are aggregated using feature\-dependent weighting mechanisms that respect covariate similarity: we implement two complementary strategies based on Nadaraya\-Watson kernel regression\[[15](https://arxiv.org/html/2606.13821#bib.bib15),[29](https://arxiv.org/html/2606.13821#bib.bib29)\]and dot\-product attention\[[12](https://arxiv.org/html/2606.13821#bib.bib12),[26](https://arxiv.org/html/2606.13821#bib.bib26)\]\. The resulting framework inherits theoretical grounding from recent advances in distributional causal inference\[[10](https://arxiv.org/html/2606.13821#bib.bib10)\]\.

Our focus on discrete dose settings, exemplified by scenarios such as the number of drugs co\-administered, addresses a practically relevant yet methodologically underexplored domain\. For a given patient, a lower dose may confer benefit while a higher dose proves detrimental, and the optimal dose likely depends on patient\-specific characteristics\. By directly modeling the probability that treatment benefit exceeds a clinically meaningful threshold \(here defined as superiority over control\), our approach provides a foundation for dose personalization grounded in probabilistic benefit estimates\.

The contributions of this work are threefold:

1. 1\.We propose a nonparametric model for IPTB estimation that avoids parametric distributional assumptions\. The attention\-based aggregation mechanism enables flexible borrowing of information from similar historical cases\. The main idea behind the model is to consider all pairs of instances such that one instance in the pair is from the control group and another instance is from the treatment group\. Considering pairs significantly increases the training sample and allows us to validate the model for real data\.
2. 2\.Our framework explicitly accommodates varying treatment doses, moving beyond the binary treatment paradigm that dominates existing IPTB literature\.
3. 3\.We evaluate the proposed approach through comprehensive numerical experiments on both real\-world and synthetic datasets, examining performance under covariate shift, varying sample sizes, and heterogeneous outcome structures\. The proposed model is compared with the well\-known meta\-learners: T\-learner, S\-learner, X\-learner, and DR\-learner\. Random forests\[[2](https://arxiv.org/html/2606.13821#bib.bib2)\]are selected as the base learners for all meta\-learners\. All codes are publicly available at[https://github\.com/NTAILab/AIPTBDose](https://github.com/NTAILab/AIPTBDose)\.

## 2Problem Formulation

This work considers an observational study comprising two distinct cohorts: a control group and a treatment group\. The control cohort is represented by the dataset𝒟0=\{\(𝐱i,yi\)\}i=1c\\mathcal\{D\}\_\{0\}=\\\{\(\\mathbf\{x\}\_\{i\},y\_\{i\}\)\\\}\_\{i=1\}^\{c\}, containingccindependent observations\. For each subjectii,𝐱i∈ℝd\\mathbf\{x\}\_\{i\}\\in\\mathbb\{R\}^\{d\}denotes add\-dimensional covariate vector, andyi∈ℝy\_\{i\}\\in\\mathbb\{R\}represents the observed outcome under the control condition \(e\.g\., survival time or a physiological measure\)\.

Similarly, the treatment cohort is denoted by𝒟1=\{\(𝐳j,hj,A\)\}j=1t\\mathcal\{D\}\_\{1\}=\\\{\(\\mathbf\{z\}\_\{j\},h\_\{j\},A\)\\\}\_\{j=1\}^\{t\}, consisting ofttsubjects\. Here,𝐳j∈ℝd\\mathbf\{z\}\_\{j\}\\in\\mathbb\{R\}^\{d\}is the covariate vector, andhj∈ℝh\_\{j\}\\in\\mathbb\{R\}is the observed outcome following the intervention\. To formalize the study design, we define a treatment dosageA∈\{1,…,m\}A\\in\\\{1,\.\.\.,m\\\}\. Generally, the parameterAAcan also be added to the control patients assuming the conditionA=0A=0\. However, this is not necessary in the context of the proposed model\.

Adopting the potential outcomes framework, letYYandHHdenote the counterfactual outcomes for a subject under control and treatment \(A≥1A\\geq 1\), respectively\. Unlike conventional approaches that target the Conditional Average Treatment Effect \(CATE\), our inference focuses on the probability of individual benefit\. Concretely, for a subject with covariate vector𝐗=𝐱\\mathbf\{X\}=\\mathbf\{x\}, we seek to estimate:

Ψ​\(𝐱\)=Pr⁡\(H\>Y∣𝐗=𝐱\)\.\\Psi\(\\mathbf\{x\}\)=\\Pr\\left\(H\>Y\\mid\\mathbf\{X\}=\\mathbf\{x\}\\right\)\.\(2\)This quantity offers a probabilistic assessment of treatment efficacy at the individual level, moving beyond an expectation of the outcome difference to capture the likelihood that treatment proves superior to control for a given patient\.

## 3Treatment Benefit Probability as a Classification Task

Three main ideas behind the proposed models are the following: \(1\) the pairwise patient comparison; \(2\) treatment doses as an additional feature in the vector of features for treatment patients; \(3\) the attention mechanism for computing probability that the difference of outcomes for a pair of patients from the control and treatment groups is positive\.

### 3\.1Pairwise patient comparison

The core methodology involves constructing pairwise comparisons between patients from distinct cohorts\. Specifically, we form pairs consisting of one treatment group patient\(𝐳i,hi\)\(\\mathbf\{z\}\_\{i\},h\_\{i\}\)and one control group patient\(𝐱j,yj\)\(\\mathbf\{x\}\_\{j\},y\_\{j\}\)\. The treatment effect for each pair is quantified asΔi​j=hi−yj\\Delta\_\{ij\}=h\_\{i\}\-y\_\{j\}\.

We collect all ordered treatment effect valuesΔi​j\\Delta\_\{ij\}fori=1,…,ti=1,\\dots,tandj=1,…,cj=1,\\dots,c, into two subsets𝒢\+\\mathcal\{G\}^\{\+\}and𝒢−\\mathcal\{G\}^\{\-\}with positive and negative valuesΔi​j\\Delta\_\{ij\}, respectively, i\.e\., we can write

𝒢\+=\{Δi​j:Δi​j\>0\},𝒢−=\{Δi​j:Δi​j≤0\}\.\\mathcal\{G\}^\{\+\}=\\\{\\Delta\_\{ij\}:\\Delta\_\{ij\}\>0\\\},\\ \\mathcal\{G\}^\{\-\}=\\\{\\Delta\_\{ij\}:\\Delta\_\{ij\}\\leq 0\\\}\.\(3\)
Let us consider subsets ofΔi​j\\Delta\_\{ij\}whose values can belong to the interval\(0,\+∞\)\(0,\+\\infty\), and introduce the index setℛ\+\\mathcal\{R\}^\{\+\}corresponding to all pairs of observations with positive valuesΔi​j\\Delta\_\{ij\}\(from𝒢\+\\mathcal\{G\}^\{\+\}\):

ℛ\+=\{\(i,j\):Δi​j\>0\}\.\\mathcal\{R\}^\{\+\}=\\\{\(i,j\):\\Delta\_\{ij\}\>0\\\}\.\(4\)
This case is simple because the probability thatΔi​j\>0\\Delta\_\{ij\}\>0is11, i\.e\.,

Pr⁡\{Δ\>0∣𝐙=𝐳i,𝐗=𝐱j\}=1,\(i,j\)∈ℛ\+\.\\Pr\\\{\\Delta\>0\\mid\\mathbf\{Z\}=\\mathbf\{z\}\_\{i\},\\mathbf\{X\}=\\mathbf\{x\}\_\{j\}\\\}=1,\\ \(i,j\)\\in\\mathcal\{R\}^\{\+\}\.\(5\)
We also introduce the index setℛ−\\mathcal\{R\}^\{\-\}such that

ℛ−=\{\(i,j\):Δi​j≤0\}\.\\mathcal\{R\}^\{\-\}=\\\{\(i,j\):\\Delta\_\{ij\}\\leq 0\\\}\.\(6\)
The corresponding cases of the differencesΔi​j=hi−yj\\Delta\_\{ij\}=h\_\{i\}\-y\_\{j\}are depicted in Fig\.[1](https://arxiv.org/html/2606.13821#S3.F1)\.

![Refer to caption](https://arxiv.org/html/2606.13821v1/Explan_IPTB_Simple.png)Figure 1:Two cases of differencesΔi​j\\Delta\_\{ij\}Our aim is find the probability that the treatment effect will be positive for a new pair of patients with feature vectors𝐳\\mathbf\{z\}and𝐱\\mathbf\{x\}, i\.e\.,

p\+​\(𝐳,𝐱\)=Pr⁡\{Δ\>0∣𝐙=𝐳,𝐗=𝐱\}\.p^\{\+\}\(\\mathbf\{z\},\\mathbf\{x\}\)=\\Pr\\\{\\Delta\>0\\mid\\mathbf\{Z\}=\\mathbf\{z\},\\mathbf\{X\}=\\mathbf\{x\}\\\}\.
In other words,p\+​\(𝐳,𝐱\)p^\{\+\}\(\\mathbf\{z\},\\mathbf\{x\}\)is a probability that the difference of observations with the treatment instance𝐳\\mathbf\{z\}and the control instance𝐱\\mathbf\{x\}producesΔi​j\\Delta\_\{ij\}which belongs to the setℛ\+\\mathcal\{R\}^\{\+\}\.

### 3\.2Treatment doses

In order to take into account different treatment doses, it is proposed to add this quantity to feature vectors of the treatment patients\. At that, there are different types of the treatment doses: discrete, continuous, and categorical\. Discrete and continuous types are represented by an additional feature whereas the categorical type is represented by a set of additional features using the one\-hot coding scheme\.

We consider the first discrete type of the treatment dosage to simplify the description\. It is important to note that the consideration of the treatment dose as an additional feature in the concatenated vectors𝐳\\mathbf\{z\}and𝐱\\mathbf\{x\}requires to weigh this feature and other features when the dimension of the feature vectors is rather large\. In this case, the weightω\\omegais assigned to the dose feature and weights\(1−ω\)/\(2​d\)\(1\-\\omega\)/\(2d\)are assigned to other features\. As a result, distances between feature vectors are calculated in accordance with the weights\.

### 3\.3Attention for solving the classification problem

Let us consider the probabilityp\+​\(𝐳,𝐱\)p^\{\+\}\(\\mathbf\{z\},\\mathbf\{x\}\)from the attention mechanism or from the kernel Nadaraya\-Watson regression point of view\. We aim to find this probability of the positive treatment effect for the pair of new patients with the feature vectors𝐳\\mathbf\{z\}and𝐱\\mathbf\{x\}\. For every pair of vectors𝐳\\mathbf\{z\}and𝐱\\mathbf\{x\}, we find the concatenations\(𝐳,𝐱\)\(\\mathbf\{z\},\\mathbf\{x\}\)and\(𝐳i,𝐱j\)\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)\. Then the Nadaraya\-Watson kernel regression can be written as follows:

p\+​\(𝐳,𝐱\)\\displaystyle p^\{\+\}\(\\mathbf\{z\},\\mathbf\{x\}\)=∑\(l,k\)∈ℛ\+∪ℛ−a​\(\(𝐳,𝐱\),\(𝐳l,𝐱k\)\)⋅𝟏​\[Δl​k\>0\]\\displaystyle=\\sum\_\{\(l,k\)\\in\\mathcal\{R\}^\{\+\}\\cup\\mathcal\{R\}^\{\-\}\}a\(\(\\mathbf\{z\},\\mathbf\{x\)\},\(\\mathbf\{z\}\_\{l\},\\mathbf\{x\}\_\{k\}\)\)\\cdot\\mathbf\{1\}\[\\Delta\_\{lk\}\>0\]=∑\(i,j\)∈ℛ\+a​\(\(𝐳,𝐱\),\(𝐳i,𝐱j\)\)⋅1\+∑\(r,s\)∈ℛ−a​\(\(𝐳,𝐱\),\(𝐳r,𝐱s\)\)⋅0,\\displaystyle=\\sum\_\{\(i,j\)\\in\\mathcal\{R\}^\{\+\}\}a\(\(\\mathbf\{z\},\\mathbf\{x\)\},\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)\)\\cdot 1\+\\sum\_\{\(r,s\)\\in\\mathcal\{R\}^\{\-\}\}a\(\(\\mathbf\{z\},\\mathbf\{x\)\},\(\\mathbf\{z\}\_\{r\},\\mathbf\{x\}\_\{s\}\)\)\\cdot 0,\(7\)where the attention weighta​\(\(𝐳,𝐱\),\(𝐳i,𝐱j\)\)a\(\(\\mathbf\{z\},\\mathbf\{x\)\},\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)\)conforms with relevance of the concatenated feature vector\(𝐳,𝐱\)\(\\mathbf\{z\},\\mathbf\{x\)\}to the concatenated feature vector\(𝐳i,𝐱j\)\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\), and there holds

∑\(r,s\)∈ℛ\+∪ℛ−a​\(\(𝐳,𝐱\),\(𝐳r,𝐱s\)\)=1\.\\sum\_\{\(r,s\)\\in\\mathcal\{R\}^\{\+\}\\cup\\mathcal\{R\}^\{\-\}\}a\(\(\\mathbf\{z\},\\mathbf\{x\)\},\(\\mathbf\{z\}\_\{r\},\\mathbf\{x\}\_\{s\}\)\)=1\.\(8\)
It can be seen from the above that Nadaraya\-Watson regression model estimatesp\+​\(𝐳,𝐱\)p^\{\+\}\(\\mathbf\{z\},\\mathbf\{x\}\)as a weighted sum of training outputs, which are indicator functions𝟏​\[Δi​j\>0\]\\mathbf\{1\}\[\\Delta\_\{ij\}\>0\], so that their weights depend on the location of\(𝐳i,𝐱j\)\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)relative to\(𝐳,𝐱\)\(\\mathbf\{z\},\\mathbf\{x\)\}\. This means that the closer\(𝐳i,𝐱j\)\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)is to\(𝐳,𝐱\)\(\\mathbf\{z\},\\mathbf\{x\)\}, the greater weight is assigned to the indicator function\. Since all pairs of indices from the setℛ\+\\mathcal\{R\}^\{\+\}satisfy the conditionΔi​j\>0\\Delta\_\{ij\}\>0, then the indicator functions are equal to 1 for all\(i,j\)∈ℛ\+\(i,j\)\\in\\mathcal\{R\}^\{\+\}as it is shown in the expression forp\+​\(𝐳,𝐱\)p^\{\+\}\(\\mathbf\{z\},\\mathbf\{x\}\)\. At the same time, the sum of attention weights for all pairs of indices\(r,s\)∈ℛ\+∪ℛ−\(r,s\)\\in\\mathcal\{R\}^\{\+\}\\cup\\mathcal\{R\}^\{\-\}is equal to 1\.

According to Nadaraya\-Watson kernel regression\[[15](https://arxiv.org/html/2606.13821#bib.bib15),[29](https://arxiv.org/html/2606.13821#bib.bib29)\], weights can be defined by means of the kernelKKas a function of the distance between the vectors𝐱i\\mathbf\{x\}\_\{i\}and𝐱\\mathbf\{x\}\. The kernel estimates how𝐱i\\mathbf\{x\}\_\{i\}is close to𝐱\\mathbf\{x\}\. Then the weight is written as follows:

a​\(\(𝐳,𝐱\),\(𝐳i,𝐱j\)\)=K​\(\(𝐳,𝐱\),\(𝐳i,𝐱j\)\)∑\(r,s\)∈ℛ\+∪ℛ−K​\(\(𝐳,𝐱\),\(𝐳r,𝐱s\)\)\.a\(\(\\mathbf\{z\},\\mathbf\{x\)\},\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)\)=\\frac\{K\(\(\\mathbf\{z\},\\mathbf\{x\)\},\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)\)\}\{\\sum\_\{\(r,s\)\\in\\mathcal\{R\}^\{\+\}\\cup\\mathcal\{R\}^\{\-\}\}K\(\(\\mathbf\{z\},\\mathbf\{x\)\},\(\\mathbf\{z\}\_\{r\},\\mathbf\{x\}\_\{s\}\)\)\}\.\(9\)
If the kernelK​\(\(𝐳,𝐱\),\(𝐳i,𝐱j\)\)K\(\(\\mathbf\{z\},\\mathbf\{x\)\},\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)\)is Gaussian, then the attention weight is expressed through the softmax operation:

a​\(\(𝐳,𝐱\),\(𝐳i,𝐱j\)\)=softmax​\(−‖\(𝐳,𝐱\)−\(𝐳i,𝐱j\)‖2τ\),a\(\(\\mathbf\{z\},\\mathbf\{x\)\},\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)\)=\\text\{\{softmax\}\}\\left\(\-\\frac\{\\left\\\|\(\\mathbf\{z\},\\mathbf\{x\)\}\-\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)\\right\\\|^\{2\}\}\{\\tau\}\\right\),\(10\)whereτ\\tauis a tuning \(temperature\) parameter\.

The numerator ina​\(\(𝐳,𝐱\),\(𝐳i,𝐱j\)\)a\(\(\\mathbf\{z\},\\mathbf\{x\)\},\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)\)measures the similarity between the query point\(𝐳,𝐱\)\(\\mathbf\{z\},\\mathbf\{x\}\)and the training pair\(𝐳i,𝐱j\)\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)\. The denominator sums over all pairs involved in the dataset, acting as a local normalizing constant\. In a special case, when the kernels are Gaussian and concatenated vectors do not include cross\-terms between𝐳\\mathbf\{z\}and𝐱\\mathbf\{x\}that prevent factorization, the attention weights are defined as:

a​\(𝐳,𝐱,𝐳i,𝐱j\)=K​\(𝐳,𝐳i\)⋅K​\(𝐱,𝐱j\)∑\(r,s\)∈ℛ\+∪ℛ−K​\(𝐳,𝐳r\)⋅K​\(𝐱,𝐱s\)\.a\(\\mathbf\{z\},\\mathbf\{x\},\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)=\\frac\{K\(\\mathbf\{z\},\\mathbf\{z\}\_\{i\}\)\\cdot K\(\\mathbf\{x\},\\mathbf\{x\}\_\{j\}\)\}\{\\sum\_\{\(r,s\)\\in\\mathcal\{R\}^\{\+\}\\cup\\mathcal\{R\}^\{\-\}\}K\(\\mathbf\{z\},\\mathbf\{z\}\_\{r\}\)\\cdot K\(\\mathbf\{x\},\\mathbf\{x\}\_\{s\}\)\}\.\(11\)
Following the attention mechanism framework\[[12](https://arxiv.org/html/2606.13821#bib.bib12),[26](https://arxiv.org/html/2606.13821#bib.bib26)\], we treat\(𝐳,𝐱\)\(\\mathbf\{z\},\\mathbf\{x\}\)as the*query*,\(𝐳i,𝐱j\)\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)as*keys*, and11as*values*\. Let us define the :

𝐪\\displaystyle\\mathbf\{q\}=𝐖Q​\(𝐳,𝐱\)⊤∈ℝd,\\displaystyle=\\mathbf\{W\}\_\{Q\}\\left\(\\mathbf\{z\},\\mathbf\{x\}\\right\)^\{\\top\}\\in\\mathbb\{R\}^\{d\},𝐤i​j\\displaystyle\\mathbf\{k\}\_\{ij\}=𝐖K​\(𝐳i,𝐱j\)⊤∈ℝd,\\displaystyle=\\mathbf\{W\}\_\{K\}\\left\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\\right\)^\{\\top\}\\in\\mathbb\{R\}^\{d\},vi​j\\displaystyle v\_\{ij\}=𝟏​\[Δi​j\>0\],\\displaystyle=\\mathbf\{1\}\[\\Delta\_\{ij\}\>0\],\(12\)where𝐖Q∈ℝd×\(2​d\+1\)\\mathbf\{W\}\_\{Q\}\\in\\mathbb\{R\}^\{d\\times\(2d\+1\)\}and𝐖K∈ℝd×\(2​d\+1\)\\mathbf\{W\}\_\{K\}\\in\\mathbb\{R\}^\{d\\times\(2d\+1\)\}are learnable weight matrices\. The attention weights become:

a​\(𝐳,𝐱,𝐳i,𝐱j\)=exp⁡\(12​d\+1​𝐪⊤​𝐤i​j\)∑\(s,r\)∈ℛ\+∪ℛ−exp⁡\(12​d\+1​𝐪⊤​𝐤s​r\)\.a\(\\mathbf\{z\},\\mathbf\{x\},\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)=\\frac\{\\exp\\left\(\\frac\{1\}\{\\sqrt\{2d\+1\}\}\\mathbf\{q\}^\{\\top\}\\mathbf\{k\}\_\{ij\}\\right\)\}\{\\sum\_\{\(s,r\)\\in\\mathcal\{R\}^\{\+\}\\cup\\mathcal\{R\}^\{\-\}\}\\exp\\left\(\\frac\{1\}\{\\sqrt\{2d\+1\}\}\\mathbf\{q\}^\{\\top\}\\mathbf\{k\}\_\{sr\}\\right\)\}\.\(13\)
Let𝐊=\[𝐤i​j\]\\mathbf\{K\}=\\left\[\\mathbf\{k\}\_\{ij\}\\right\]contain all keys and𝐕=\[vi​j\]\\mathbf\{V\}=\\left\[v\_\{ij\}\\right\]contain all values\. The distribution can be expressed in the matrix form:

p\+​\(𝐳,𝐱\)=softmax​\(12​d\+1​𝐪𝐊⊤\)​𝐕\.p^\{\+\}\(\\mathbf\{z\},\\mathbf\{x\}\)=\\text\{softmax\}\\left\(\\frac\{1\}\{\\sqrt\{2d\+1\}\}\\mathbf\{q\}\\mathbf\{K\}^\{\\top\}\\right\)\\mathbf\{V\}\.\(14\)
When features have different importance, we introduce a diagonal weight matrix𝐖∈ℝd×\(2​d\+1\)\\mathbf\{W\}\\in\\mathbb\{R\}^\{d\\times\(2d\+1\)\}:

𝐖=diag⁡\(ω1,ω2,…,ω2​d\+1\),\\mathbf\{W\}=\\operatorname\{diag\}\(\\omega\_\{1\},\\omega\_\{2\},\\dots,\\omega\_\{2d\+1\}\),\(15\)whereωi≥0\\omega\_\{i\}\\geq 0represents the importance of theii\-th feature dimension\. Hence, the weighted queries and weighted keys are𝐐~=𝐐𝐖1/2\\widetilde\{\\mathbf\{Q\}\}=\\mathbf\{Q\}\\mathbf\{W\}^\{1/2\}and𝐊~=𝐊𝐖1/2\\widetilde\{\\mathbf\{K\}\}=\\mathbf\{K\}\\mathbf\{W\}^\{1/2\}, respectively, where𝐖1/2=diag⁡\(ω1,…,ω2​d\+1\)\\mathbf\{W\}^\{1/2\}=\\operatorname\{diag\}\(\\sqrt\{\\omega\_\{1\}\},\\dots,\\sqrt\{\\omega\_\{2d\+1\}\}\)\.

The model learns matrices𝐖Q\\mathbf\{W\}\_\{Q\}and𝐖K\\mathbf\{W\}\_\{K\}by minimizing the log\-likelihood loss with theL2L\_\{2\}\-regularization:

ℒ\(𝐩\\displaystyle\\mathcal\{L\}\(\\mathbf\{p\}∣𝐖Q,𝐖K\)=−∑\(i,j\)∈ℛ\+log\(p\+\(𝐳i,𝐱j\)\)−∑\(i,j\)∈ℛ−log\(p−\(𝐳i,𝐱j\)\)\\displaystyle\\mid\\mathbf\{W\}\_\{Q\},\\mathbf\{W\}\_\{K\}\)=\-\\sum\_\{\(i,j\)\\in\\mathcal\{R\}^\{\+\}\}\\log\\left\(p^\{\+\}\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)\\right\)\-\\sum\_\{\(i,j\)\\in\\mathcal\{R\}^\{\-\}\}\\log\\left\(p^\{\-\}\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)\\right\)\+η​\(‖𝐖Q‖2\+‖𝐖K‖2\),\\displaystyle\+\\eta\\left\(\\\|\\mathbf\{W\}\_\{Q\}\\\|^\{2\}\+\\\|\\mathbf\{W\}\_\{K\}\\\|^\{2\}\\right\),\(16\)whereη\\etais the hyperparameter controlling the strength of theL2L\_\{2\}\-regularization\.

During inference, the probabilityp\+​\(𝐱,𝐱\)p^\{\+\}\(\\mathbf\{x\},\\mathbf\{x\}\)is computed for a new instance𝐱\\mathbf\{x\}under conditionsA=1,…,mA=1,\.\.\.,m\. Suppose thatγ\\gammais the threshold of the probabilityp\+​\(𝐱,𝐱\)p^\{\+\}\(\\mathbf\{x\},\\mathbf\{x\}\)for decision making about the positive treatment effect\. Then the number of doses is selected minimizing the difference betweenp\+​\(𝐱,𝐱\)p^\{\+\}\(\\mathbf\{x\},\\mathbf\{x\}\)andγ\\gammabyp\+​\(𝐱,𝐱\)≥γp^\{\+\}\(\\mathbf\{x\},\\mathbf\{x\}\)\\geq\\gamma\.

The quantityp\+​\(𝐳,𝐱\)p^\{\+\}\(\\mathbf\{z\},\\mathbf\{x\}\)can be viewed as a*causal risk*function reflecting the probability that the treatment yields a beneficial outcome for an individual with covariates𝐱\\mathbf\{x\}\. Direct estimation ofΔ\\Deltaat the individual level is infeasible because, for each individual, only one potential outcome \(either under treatment or control\) is observed\. The proposed approach employs a*contrastive*approach inspired by classification: by forming pairs\(𝐳i,hi\)\(\\mathbf\{z\}\_\{i\},h\_\{i\}\)and\(𝐱j,yj\)\(\\mathbf\{x\}\_\{j\},y\_\{j\}\)from treatment and control groups, respectively, we approximate the distribution of individual effects indirectly\. The sign of the observed differencehi−yjh\_\{i\}\-y\_\{j\}for these pairs acts as a surrogate indicator of the unobserved individual effect, assuming that individuals with similar covariates have similar potential outcomes\. This assumption aligns with the*conditional exchangeability*condition common in causal inference\.

By aggregating large numbers of such pairs and weighting their contributions according to covariate similarity \(via attention or kernel methods\), the model effectively estimates the probability that a new individual with covariates𝐱\\mathbf\{x\}would benefit from treatment instances𝐳\\mathbf\{z\}\. The likelihood thatΔ\>0\\Delta\>0for a given pair is thus approximated by the label indicating positive outcomes of their difference, up to the approximation quality provided by the weighting scheme\.

## 4Numerical Experiments

To assess model performance, we divide all experiments into two parts\. The first part of numerical experiments can be conducted on both synthetic and real data\. Predicted probabilitiesPr⁡\{Δ\>0∣𝐙=𝐳i,𝐗=𝐱j\}\\Pr\\left\\\{\{\\Delta\>0\\mid\\mathbf\{Z\}=\\mathbf\{z\}\_\{i\},\\mathbf\{X\}=\\mathbf\{x\}\_\{j\}\}\\right\\\}are calculated for each pair\(𝐳i,𝐱j\)\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\)from the testing set\. In this case, we do not need to consider only pairs\(𝐱,𝐱\)\(\\mathbf\{x\},\\mathbf\{x\}\)because Dose\-AIPTB allows us to estimatep\+​\(𝐳,𝐱\)p^\{\+\}\(\\mathbf\{z\},\\mathbf\{x\}\)for arbitrary pairs composed from the treatment and control groups\. The corresponding validation metrics will be called Val 1\. The second part is implemented using only synthetic datasets\. It estimates IPTB for patients whose feature vectors are identical in both the treatment and control groups\. The corresponding outcomeshih\_\{i\}andyjy\_\{j\}are generated in accordance with the dataset rules\. After generating control and treatment outcomes for many feature vectors, for each pair\(𝐱,𝐱\)\(\\mathbf\{x\},\\mathbf\{x\}\), we can calculate the predicted probabilityp\+​\(𝐱,𝐱\)=Pr⁡\{Δ\>0∣𝐙=𝐱,𝐗=𝐱\}p^\{\+\}\(\\mathbf\{x\},\\mathbf\{x\}\)=\\Pr\\left\\\{\{\\Delta\>0\\mid\\mathbf\{Z\}=\\mathbf\{x\},\\mathbf\{X\}=\\mathbf\{x\}\}\\right\\\}using the proposed model\. The corresponding validation metrics will be called Val 2\.

For quantitative evaluation, we provide the area under the receiver operating characteristic curve \(ROC\-AUC\)\. Each metric value is obtained by averaging results from five\-fold stratified cross\-validation repeated ten times with different random seeds to ensure statistical reliability\. For every dataset, we provide two types of graphs with the ROC\-curves\. The first type \(left graphs\) illustrates the training, validation\-1 \(Val 1\) and validation\-2 \(Val 2\) ROC\-curves for Dose\-AIPTB, where Val 1 and Val 2 correspond to validation on pairs\(𝐳,𝐱\)\(\\mathbf\{z\},\\mathbf\{x\}\)and\(𝐱,𝐱\)\(\\mathbf\{x\},\\mathbf\{x\}\), respectively\.

Dose\-AIPTB is compared with meta\-learners: the T\-learner, S\-learner, X\-learner, and DR\-learner\. Each meta\-learner is trained using the random forest regressor consisting of 200 random trees\. It should be noted that the meta\-learners return CATE values, so the sign of the CATE is converted into the IPTB class label\. At that, meta\-learners are trained on every dose\.

### 4\.1Synthetic data

We study the proposed model Dose\-AIPTB on synthetic datasets:*Simple, Linear, Step\-wise, Spiral, Power, Weibull*\. They are generated in accordance with the following functions of the same name\.

1. 1\.*Simple*: Instances are generated such that y​\(A\)=\(1−0\.25⋅A\)⋅x\(1\)\+\(0\.25\+0\.25⋅A\)⋅x\(2\)\.y\(A\)=\(1\-0\.25\\cdot A\)\\cdot x^\{\(1\)\}\+\(0\.25\+0\.25\\cdot A\)\\cdot x^\{\(2\)\}\.\(17\) HereAAis uniformly generated from\{1,2,3\}\\\{1,2,3\\\}for treatments andA=0A=0for controls\. Numbers of controls and treatments arec=350c=350andt=50t=50, respectively\.
2. 2\.*Linear*: Feature valuesx\(i\)x^\{\(i\)\}are uniformly generated from\[0,1\]\[0,1\], i\.e\.,x\(i\)∼𝒰​\(0,1\)x^\{\(i\)\}\\sim\\mathcal\{U\}\(0,1\), values ofyyare generated using the linear function: y​\(A\)=2⋅\(1−0\.2⋅A\)⋅x\(1\)\+4⋅x\(2\)\+8⋅0\.2⋅A⋅x\(3\)\.y\(A\)=2\\cdot\(1\-0\.2\\cdot A\)\\cdot x^\{\(1\)\}\+4\\cdot x^\{\(2\)\}\+8\\cdot 0\.2\\cdot A\\cdot x^\{\(3\)\}\.\(18\) Numbers of controls and treatments are the same as in the Simple dataset;A∈\{0,1,2,3,4,5\}A\\in\\\{0,1,2,3,4,5\\\}\.
3. 3\.*Step\-wise*: Instances are generated such that y​\(A\)\\displaystyle y\(A\)=\(1−0\.25⋅A⋅𝕀​\(x\(1\)<0\.5\)\)⋅x\(1\)\\displaystyle=\(1\-0\.25\\cdot A\\cdot\\mathbb\{I\}\\left\(x^\{\(1\)\}<0\.5\\right\)\)\\cdot x^\{\(1\)\}\+\(0\.25\+0\.25⋅A\)⋅x\(2\)⋅𝕀​\(x\(2\)<0\.5\)\.\\displaystyle\+\(0\.25\+0\.25\\cdot A\)\\cdot x^\{\(2\)\}\\cdot\\mathbb\{I\}\\left\(x^\{\(2\)\}<0\.5\\right\)\.\(19\) HereA∈\{0,1,2,3\}A\\in\\\{0,1,2,3\\\}\. The dosage increases the weight ofx\(2\)x^\{\(2\)\}in the outcome while decreasing the weight ofx\(1\)x^\{\(1\)\}\. The treatment effect in this dataset can be either positive or negative\. Parameters of the Step\-wise dataset are the same as for the Simple dataset, but whenx\(1\)\>0\.5x^\{\(1\)\}\>0\.5, its coefficient becomes11, and whenx\(2\)\>0\.5x^\{\(2\)\}\>0\.5, its coefficient is set to0\. This rule forx\(1\)\>0\.5x^\{\(1\)\}\>0\.5removes the negative effect of the dosage, while forx\(2\)\>0\.5x^\{\(2\)\}\>0\.5, it removes the positive effect\. When bothx\(1\)\>0\.5x^\{\(1\)\}\>0\.5andx\(2\)\>0\.5x^\{\(2\)\}\>0\.5, the dosage becomes irrelevant \(yydepends linearly only onx\(1\)x^\{\(1\)\}\)\.
4. 4\.*Spiral*: Vectors𝐱∈ℝ5\\mathbf\{x\}\\in\\mathbb\{R\}^\{5\}are generated by using the Archimedean spiral as follows: 𝐱=\(t​sin⁡\(t\),t​cos⁡\(t\),…,t​sin⁡\(t⋅d/2\),t​cos⁡\(t⋅d/2\)\),\\mathbf\{x\}=\(t\\sin\(t\),t\\cos\(t\),\.\.\.,t\\sin\(t\\cdot d/2\),t\\cos\(t\\cdot d/2\)\),\(20\)for evendd, and 𝐱=\(t​sin⁡\(t\),t​cos⁡\(t\),…,t​sin⁡\(t⋅⌈d/2⌉\)\)\.\\mathbf\{x\}=\(t\\sin\(t\),t\\cos\(t\),\.\.\.,t\\sin\(t\\cdot\\left\\lceil d/2\\right\\rceil\)\)\.\(21\)for odddd\. Values ofyyare generated as y​\(A\)=\(a\+0\.2⋅A\)⋅t−1\.5⋅A⋅b,y\(A\)=\(a\+0\.2\\cdot A\)\\cdot t\-1\.5\\cdot A\\cdot b,\(22\)wherea∼𝒰​\(0\.6,1\)a\\sim\\mathcal\{U\}\(0\.6,1\),b∼𝒰​\(0\.6,1\)b\\sim\\mathcal\{U\}\(0\.6,1\),t∼𝒰​\(1,12\)t\\sim\\mathcal\{U\}\(1,12\),d=5d=5,AAis uniformly generated from\{1,2,3,4,5\}\\\{1,2,3,4,5\\\}\. Ifttis large, increasingAAwill increasey​\(A\)y\(A\), otherwise, it will decreasey​\(A\)y\(A\)\. Herec=400c=400andt=200t=200\.
5. 5\.*Power*:Feature vectors𝐱∈ℝ5\\mathbf\{x\}\\in\\mathbb\{R\}^\{5\}are generated by using the following representation: 𝐱=\(t1/d,t2/d,…,td/d\)\.\\mathbf\{x\}=\(t^\{1/\\sqrt\{d\}\},t^\{2/\\sqrt\{d\}\},\.\.\.,t^\{d/\\sqrt\{d\}\}\)\.\(23\)Outcomes are computed as: y​\(A\)=\(a−0\.2⋅A\)⋅exp⁡\(−\(t−s\)2b−0\.1⋅\(A−5\)\)\.y\(A\)=\(a\-0\.2\\cdot A\)\\cdot\\exp\\left\(\-\\frac\{\(t\-s\)^\{2\}\}\{b\-0\.1\\cdot\(A\-5\)\}\\right\)\.\(24\) Herea∼𝒰​\(9,10\)a\\sim\\mathcal\{U\}\(9,10\),b∼𝒰​\(0\.5,1\)b\\sim\\mathcal\{U\}\(0\.5,1\),t∼𝒰​\(0,7\)t\\sim\\mathcal\{U\}\(0,7\),s=3\.5s=3\.5,d=5d=5,AAis uniformly generated from\{1,2,3,4,5\}\\\{1,2,3,4,5\\\}\. The effect ofAAony​\(A\)y\(A\)is positive when\(t−s\)\(t\-s\)is large and negative when\(t−s\)\(t\-s\)is small\.
6. 6\.*Weibull*: Each observation has two features,x\(1\)x^\{\(1\)\}andx\(2\)x^\{\(2\)\}sampled uniformly from0to11\. The event time \(outcome\) is generated from the Weibull distribution with shape parameterk=3\.0k=3\.0as follows: T=y​\(A\)Γ​\(1\+1k\)⋅\(−log⁡\(u\)\)1k\.T=\\frac\{y\(A\)\}\{\\Gamma\\left\(1\+\\frac\{1\}\{k\}\\right\)\}\\cdot\\left\(\-\\log\(u\)\\right\)^\{\\frac\{1\}\{k\}\}\.\(25\)HereΓ​\(⋅\)\\Gamma\(\\cdot\)is the gamma function;u∼𝒰​\(0,1\)u\\sim\\mathcal\{U\}\(0,1\);y​\(A\)y\(A\)is determined as in the Simple dataset;A∈\{0,1,2,3\}A\\in\\\{0,1,2,3\\\}\.

Fig\.[2](https://arxiv.org/html/2606.13821#S4.F2)displays two side\-by\-side Receiver Operating Characteristic \(ROC\) plots obtained for the*Simple*dataset\. The left plot shows the performance of the proposed model Dose\-AIPTB on training and validation sets, where all curves \(Train, Val 1, Val 2\) are tightly clustered near the top left with high AUC scores ranging from roughly0\.9850\.985to0\.9980\.998\. The right plot compares different meta\-learner algorithms, showing that the Dose\-AIPTB model \(red solid line\) achieves the best performance with an AUC of0\.99810\.9981, significantly outperforming the T\-learner, S\-learner, X\-learner, and DR\-learner\.

![Refer to caption](https://arxiv.org/html/2606.13821v1/Roc_curves_aggregated_Simple.png)Figure 2:Left plot: The ROC curves and ROC\-AUC scores obtained on the training, Val 1, and Val 2 sets for Dose\-AIPTB trained on the Simple dataset\. Right plot: Comparison of ROC curves and ROC\-AUC scores for Dose\-AIPTB and the meta\-learners\.Fig\.[3](https://arxiv.org/html/2606.13821#S4.F3)illustrates values of training \(Train\) and two testing \(Val 1 and Val 2\) loss functions depending on the epoch numbers for the Simple dataset\. One can see from the figure that two curves \(Train and Val 1\) follow an almost identical path\. They drop rapidly until about epoch 40\. Afterward, they flatten out, slowly decreasing to a final loss of around0\.10\.1\. The third curve \(Val 2\) behaves differently\. It is decreasing to a final loss of around0\.20\.2\.

![Refer to caption](https://arxiv.org/html/2606.13821v1/Loss_aggregated_Simple.png)Figure 3:Training and validation loss functions for the Simple datasetFig\.[4](https://arxiv.org/html/2606.13821#S4.F4)displays similar ROC plots obtained for the*Linear*dataset\. The left plot shows the performance of the proposed model Dose\-AIPTB on training and validation sets\. The right plot also demonstrates that the Dose\-AIPTB model \(red solid line\) achieves the best performance with an AUC of0\.95440\.9544, significantly outperforming the T\-learner, S\-learner, X\-learner, and DR\-learner\.

![Refer to caption](https://arxiv.org/html/2606.13821v1/Roc_curves_aggregated_Linear.png)Figure 4:Left plot: The ROC curves and ROC\-AUC scores obtained on the training, Val 1, and Val 2 sets for Dose\-AIPTB trained on the Linear dataset\. Right plot: Comparison of ROC curves and ROC\-AUC scores for Dose\-AIPTB and the meta\-learners\.Fig\.[5](https://arxiv.org/html/2606.13821#S4.F5)illustrates values of training \(Train\) and two testing \(Val 1 and Val 2\) loss functions depending on the epoch numbers for the Linear dataset\. Two curves \(Train and Val 1\) also follow an almost identical path\. The third curve \(Val 2\) starts lower than the others and plateaus at a significantly higher loss level compared to the Train and Val 1 functions suggesting the model performs worse on this specific validation set\.

![Refer to caption](https://arxiv.org/html/2606.13821v1/Loss_aggregated_Linear.png)Figure 5:Training and validation loss functions for the Linear datasetA more complex dataset is*Step\-wise*\. Similar results for this dataset under condition of training on 350 controls and 50 treatments are shown in Fig\.[6](https://arxiv.org/html/2606.13821#S4.F6)\. It can be seen from the figure that Dose\-AIPTB is comparable with the X\-learner\. It is interesting to point out that results change if the models are trained on 50 controls and 350 treatments as shown in Fig\.[7](https://arxiv.org/html/2606.13821#S4.F7)\.

![Refer to caption](https://arxiv.org/html/2606.13821v1/Roc_curves_aggregated_Step_350c_50t.png)Figure 6:Left plot: The ROC curves and ROC\-AUC scores obtained on the training, Val 1, and Val 2 sets for Dose\-AIPTB trained on the Step\-wise dataset consisting of 350 controls and 50 treatments\. Right plot: Comparison of ROC curves and ROC\-AUC scores for Dose\-AIPTB and the meta\-learners\.![Refer to caption](https://arxiv.org/html/2606.13821v1/Roc_curves_aggregated_Step_50c_350t.png)Figure 7:ROC\-AUC scores obtained on the training, Val 1, and Val 2 sets for Dose\-AIPTB trained on the Step\-wise dataset consisting of 50 controls and 350 treatments\. Right plot: Comparison of ROC curves and ROC\-AUC scores for Dose\-AIPTB and the meta\-learners\.Figs\.[8](https://arxiv.org/html/2606.13821#S4.F8)and[9](https://arxiv.org/html/2606.13821#S4.F9)illustrate results for the*Spiral*and*Power*datasets, respectively\.

![Refer to caption](https://arxiv.org/html/2606.13821v1/Roc_curves_aggregated_Spiral.png)Figure 8:Left plot: The ROC curves and ROC\-AUC scores obtained on the training, Val 1, and Val 2 sets for Dose\-AIPTB trained on the Spiral dataset\. Right plot: Comparison of ROC curves and ROC\-AUC scores for Dose\-AIPTB and the meta\-learners\.![Refer to caption](https://arxiv.org/html/2606.13821v1/Roc_curves_aggregated_Power.png)Figure 9:Left plot: The ROC curves and ROC\-AUC scores obtained on the training, Val 1, and Val 2 sets for Dose\-AIPTB trained on the Power dataset\. Right plot: Comparison of ROC curves and ROC\-AUC scores for Dose\-AIPTB and the meta\-learners\.Fig\.[10](https://arxiv.org/html/2606.13821#S4.F10)shows results for the*Weibull*datasets generated with different numbers of controls and treatments\. It can be seen from the plots that Dose\-AIPTB does not change its performance in contrast to other models\. The extended results are shown in Table[1](https://arxiv.org/html/2606.13821#S4.T1)\. This table presents AUC values with standard deviations for five models trained on the Weibull dataset under varying ratios of controls \(c\) to treatments \(t\)\. The Dose\-AIPTB method consistently achieves the highest performance across all seven experimental settings, maintaining AUC scores between0\.7240\.724and0\.7370\.737regardless of the sample split\. In contrast, the meta\-learner baselines \(T\-learner, S\-learner, X\-learner, and DR\-learner\) exhibit lower and more variable performance, with AUC values generally ranging from0\.5750\.575to0\.6750\.675\. Notably, the performance of Dose\-AIPTB remains remarkably stable even when the treatment group becomes very small \(e\.g\., 50 treatments vs\. 350 controls\), whereas other models tend to degrade more significantly under imbalanced conditions\. The most challenging configuration appears to be the extreme imbalance of 350 controls to 50 treatments, where all models except Dose\-AIPTB drop to their lowest or near\-lowest scores\.

![Refer to caption](https://arxiv.org/html/2606.13821v1/Roc_curves_aggregated_Weibull.png)Figure 10:Left plots: The ROC curves and ROC\-AUC scores obtained on the training, Val 1, and Val 2 sets for Dose\-AIPTB trained on the Weibull datasets which differ by numbers of controls and treatments\. Right plots: The corresponding comparison of ROC curves and ROC\-AUC scores for Dose\-AIPTB and the meta\-learners\.Table 1:AUC values for the models trained on the Weibull dataset by different numbers of controls \(c\) and treatments \(t\)
### 4\.2Real data

Our empirical evaluation leverages the widely adopted Infant Health and Development Program \(IHDP\) dataset which can be viewed as a benchmark resource frequently employed for heterogeneous treatment effect \(HTE\) estimation\[[5](https://arxiv.org/html/2606.13821#bib.bib5)\]\. Originally compiled to assess how specialist\-conducted home visits influence later cognitive outcomes in preterm infants, the dataset comprises 747 participants characterized by 25 covariates: 6 continuous and 19 binary variables capturing key attributes of both infants and their mothers\. Notably, the experimental setup involves 139 distinct treatment configurations\. The dataset is publicly available via the repository at[https://github\.com/vdorie/npci](https://github.com/vdorie/npci), facilitating reproducibility and comparative analysis across causal inference methodologies\. Due to the large dimension of the instances in the dataset, we introduce weights of features such that the feature with index2​d\+12d\+1\(the dose feature\) has the weightω2​d\+1=0\.3\\omega\_\{2d\+1\}=0\.3, other features have weightsωi=0\.7/\(2​d\)\\omega\_\{i\}=0\.7/\(2d\),i=1,…,di=1,\.\.\.,d\.

Fig\.[11](https://arxiv.org/html/2606.13821#S4.F11)displays ROC plots obtained for the*IHDP*dataset\. The left plot again shows the performance of the proposed model Dose\-AIPTB\. The right plot compares different meta\-learner algorithms, showing that the Dose\-AIPTB model achieves the best performance with an AUC of0\.98100\.9810, significantly outperforming the T\-learner, S\-learner, X\-learner, and DR\-learner\. Fig\.[12](https://arxiv.org/html/2606.13821#S4.F12)illustrates the loss functions depending on the epoch numbers for the IHDP dataset\. All three curves show a steep decline in loss initially, dropping from above0\.450\.45to below0\.300\.30within the first 20 epochs, before the rate of improvement slows significantly\. By the end of the training, the training set reaches the lowest loss near0\.200\.20, whereas the validation sets plateau at slightly higher values, with Val 2 maintaining the highest loss of roughly0\.230\.23\.

![Refer to caption](https://arxiv.org/html/2606.13821v1/Roc_curves_aggregated_IHDP100.png)Figure 11:Left plot: The ROC curves and ROC\-AUC scores obtained on the training, Val 1, and Val 2 sets for Dose\-AIPTB trained on the IHDP dataset\. Right plot: Comparison of ROC curves and ROC\-AUC scores for Dose\-AIPTB and the meta\-learners\.![Refer to caption](https://arxiv.org/html/2606.13821v1/Loss_aggregated_IHDP.png)Figure 12:Training and validation loss functions for the IHDP datasetTable[2](https://arxiv.org/html/2606.13821#S4.T2)provides AUC values for all considered datasets obtained by Dose\-AIPTB and all studied meta\-models\. This table presents the AUC values, including standard deviations, for five different models trained across seven distinct datasets\. The models being compared are Dose\-AIPTB, T\-learner, S\-learner, X\-learner, and DR\-learner\. Dose\-AIPTB demonstrates superior performance in the majority of the scenarios\. However, on the Spiral dataset, the S\-learner slightly outperforms Dose\-AIPTB with an AUC of0\.8490\.849compared to0\.8430\.843\. The Weibull dataset appears to be the most challenging for all algorithms, resulting in the lowest overall AUC values which range from roughly0\.5750\.575to0\.7240\.724\.

Table 2:AUC values for the models trained on different datasets

## 5Conclusion

We introduced Dose\-AIPTB, a nonparametric framework for estimating the Individual Probability of Treatment Benefit \(IPTB\) when treatments are administered at discrete dose levels\. By reformulating IPTB estimation as a binary classification task based on pairwise patient comparisons, our approach directly addresses the gap between population\-level causal inference and individualized clinical decision\-making\. The core idea lies in an attention\-based aggregation mechanism that leverages similarity\-weighted comparisons to construct probabilistic benefit estimates while naturally incorporating dose information\. An important property of Dose\-AIPTB is that it can be validated on real data\. The corresponding testing set \(Validation 1\) is composed of pairs formed from instances in the control and treatment groups\. Since we know the values ofΔi​j\\Delta\_\{ij\}for every pair in the testing set, as well as the probabilityp\+​\(𝐳i,𝐱j\)p^\{\+\}\(\\mathbf\{z\}\_\{i\},\\mathbf\{x\}\_\{j\}\), we can estimate the model performance\.

Comprehensive experiments across six synthetic datasets and the well\-known IHDP benchmark demonstrate that Dose\-AIPTB consistently outperforms established meta\-learners \(T\-learner, S\-learner, X\-learner, DR\-learner\) in ROC\-AUC performance\. Notably, Dose\-AIPTB exhibits remarkable stability under severe sample imbalance, maintaining consistent performance where baseline methods degrade substantially\. The method’s only marginal underperformance occurred on the Spiral dataset, suggesting that certain complex non\-linear response surfaces remain challenging for pairwise comparison frameworks\. Key limitations include the current focus on discrete rather than continuous doses, computational complexity scaling with pairwise comparisons, and reliance on standard causal assumptions\. Future work will extend the framework to continuous dose spaces, integrate uncertainty quantification via conformal prediction, and develop scalable approximations for large\-scale applications\. Applying Dose\-AIPTB to high\-dimensional covariates and evaluating its clinical utility through prospective studies represent important translational next steps\. By shifting inference from expected effects to probabilities of individual benefit while explicitly modeling dose heterogeneity, Dose\-AIPTB offers a principled foundation for dose personalization in precision medicine with code publicly available at[https://github\.com/NTAILab/AIPTBDose](https://github.com/NTAILab/AIPTBDose)\.

\{credits\}

#### 5\.0\.1Acknowledgements

This work is supported by the Russian Science Foundation under grant 25\-11\-00021\.

## References

- \[1\]Béal, J\., Latouche, A\.: Causal inference with multiple versions of treatment and application to personalized medicine \(May 2020\), arXiv:2005\.12427
- \[2\]Breiman, L\.: Random forests\. Machine learning45\(1\), 5–32 \(2001\)
- \[3\]Curth, A\., van der Schaar, M\.: Nonparametric estimation of heterogeneous treatment effects: From theory to learning algorithms\. In: International Conference on Artificial Intelligence and Statistics\. pp\. 1810–1818\. PMLR \(2021\)
- \[4\]Guo, Z\., Zheng, S\., Liu, Z\., Yan, K\., Zhu, Z\.: Cetransformer: Casual effect estimation via transformer based representation learning\. In: Pattern Recognition and Computer Vision\. PRCV 2021\. Lecture Notes in Computer Science, vol\. 13022, pp\. 524–535\. Springer, Cham \(2021\)
- \[5\]Hill, J\.: Bayesian nonparametric modeling for causal inference\. Journal of Computational and Graphical Statistics20\(1\), 217–240 \(2011\)
- \[6\]Hızlı, C\., John, S\., Juuti, A\., Saarinen, T\., Pietiläinen, K\., Marttinen, P\.: Causal modeling of policy interventions from treatment–outcome sequences\. In: Proceedings of the 40th International Conference on Machine Learning\. pp\. 13050–13084\. No\. 530 in ICML’23, JMLR\.org \(2024\)
- \[7\]Imbens, G\.: Nonparametric estimation of average treatment effects under exogeneity: A review\. Review of Economics and Statistics86\(1\), 4–29 \(2004\)
- \[8\]Jeng, X\., Lu, W\., Peng, H\.: High\-dimensional inference for personalized treatment decision\. Electronic Journal of Statistics12, 12 2074–2089 \(2018\)
- \[9\]Kennedy, E\.H\.: Towards optimal doubly robust estimation of heterogeneous causal effects\. Electronic Journal of Statistics17\(2\), 3008–3049 \(2023\)
- \[10\]Konstantinov, A\., Utkin, L\., Efremenko, V\., Muliukha, V\., Lukashin, A\., Verbova, N\.: Survival analysis as imprecise classification with trainable kernels\. Mathematics13\(18\), 3040 \(2025\)
- \[11\]Kunzel, S\., Sekhon, J\., Bickel, P\., Yu, B\.: Metalearners for estimating heterogeneous treatment effects using machine learning\. PNAS116\(10\), 4156–4165 \(2019\)
- \[12\]Luong, T\., Pham, H\., Manning, C\.: Effective approaches to attention\-based neural machine translation\. In: Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing\. pp\. 1412–1421\. Association for Computational Linguistics, Lisbon, Portugal \(2015\)
- \[13\]Melnychuk, V\., Feuerriegel, S\., van der Schaar, M\.: Quantifying aleatoric uncertainty of the treatment effect: a novel orthogonal learner\. In: Advances in Neural Information Processing Systems\. vol\. 37, pp\. 105039–105089\. Curran Associates, Inc\. \(2024\)
- \[14\]Melnychuk, V\., Frauen, D\., Feuerriegel, S\.: Causal transformer for estimating counterfactual outcomes\. In: International conference on machine learning\. pp\. 15293–15329\. PMLR \(2022\)
- \[15\]Nadaraya, E\.: On estimating regression\. Theory of Probability & Its Applications9\(1\), 141–142 \(1964\)
- \[16\]Nair, N\., Gurumoorthy, K\., Mandalapu, D\.: Individual treatment effect estimation through controlled neural network training in two stages \(Jan 2022\), arXiv:2201\.08559
- \[17\]Parikh, H\., Lanners, Q\., Akras, Z\., Zafar, S\., Westover, M\., Rudin, C\., Volfovsky, A\.: Safe and interpretable estimation of optimal treatment regimes\. In: International Conference on Artificial Intelligence and Statistics\. pp\. 2134–2142\. PMLR \(2024\)
- \[18\]Park, J\., Shalit, U\., Scholkopf, B\., Muandet, K\.: Conditional distributional treatment effect with kernel conditional mean embeddings and u\-statistic regression\. In: Proceedings of the 38 th International Conference on Machine Learning, PMLR\. vol\. 139, pp\. 8401–8412 \(2021\)
- \[19\]Piskorz, J\., Kacprzyk, K\., Amad, H\., van der Schaar, M\.: Beyond the ate: Interpretable modelling of treatment effects over dose and time \(Jul 2025\), arXiv:2507\.07271
- \[20\]Qin, T\., Wang, T\.Z\., Zhou, Z\.H\.: Budgeted heterogeneous treatment effect estimation\. In: Proceedings of the 38th International Conference on Machine Learning, PMLR\. vol\. 139, pp\. 8693–8702 \(2021\)
- \[21\]Salditt, M\., Eckes, T\., Nestler, S\.: A tutorial introduction to heterogeneous treatment effect estimation with meta\-learners\. Administration and Policy in Mental Health and Mental Health Services Research51\(5\), 650–673 \(2024\)
- \[22\]Schröder, M\., Frauen, D\., Schweisthal, J\., Hess, K\., Melnychuk, V\., Feuerriegel, S\.: Conformal prediction for causal effects of continuous treatments\. In: The Thirty\-ninth Annual Conference on Neural Information Processing Systems \(2025\),[https://openreview\.net/forum?id=1nL84tQNnK](https://openreview.net/forum?id=1nL84tQNnK)
- \[23\]Schwab, P\., Linhardt, L\., Bauer, S\., Buhmann, J\., Karlen, W\.: Learning counterfactual representations for estimating individual dose\-response curves\. In: Proceedings of the AAAI Conference on Artificial Intelligence\. vol\. 34, pp\. 5612–5619 \(2020\)
- \[24\]Shi, P\., Zhong, W\., Zhang, X\., Wang, N\., Fu, X\., Wang, W\., Jin, Y\.: Estimating conditional average treatment effects via sufficient representation learning\. In: Proceedings of the Thirty\-Third International Joint Conference on Artificial Intelligence\. pp\. 4894–4901 \(2024\)
- \[25\]VanderWeele, T\., Hernan, M\.: Causal inference under multiple versions of treatment\. Journal of Causal Inference1\(1\), 1–20 \(2013\)
- \[26\]Vaswani, A\., Shazeer, N\., Parmar, N\., Uszkoreit, J\., Jones, L\., Gomez, A\., Kaiser, L\., Polosukhin, I\.: Attention is all you need\. In: Advances in Neural Information Processing Systems\. pp\. 5998–6008 \(2017\)
- \[27\]Wang, Y\., Wu, P\., Liu, Y\., Weng, C\., Zeng, D\.: Learning optimal individualized treatment rules from electronic health record data\. In: IEEE International Conference on Healthcare Informatics \(ICHI\)\. pp\. 65–71\. IEEE \(2016\)
- \[28\]Wang, Y\., Li, H\., Zhu, M\., Wu, A\., Li, B\., Yin, K\., Xiong, R\., Wu, F\., Kuang, K\.: Causal inference with complex treatments: A survey\. ACM Computing Surveys58\(9\), 1–36 \(2026\)
- \[29\]Watson, G\.: Smooth regression analysis\. Sankhya: The Indian Journal of Statistics, Series A pp\. 359–372 \(1964\)
- \[30\]Zhang, Y\.F\., Zhang, H\., Lipton, Z\., Li, L\.E\., Xing, E\.P\.: Can transformers be strong treatment effect estimators? \(Feb 2022\), arXiv:2202\.01336
- \[31\]Zhang, Y\.F\., Zhang, H\., Lipton, Z\., Li, L\.E\., Xing, E\.P\.: Exploring transformer backbones for heterogeneous treatment effect estimation \(May 2022\), arXiv:2202\.01336