DoseSurv: Predicting Personalized Survival Outcomes under Continuous-Valued Treatments
Authors: Moritz Gögl
Yu Liu
Christopher Yau
Peter Watkinson
Tingting Zhu
Venue: NeurIPS 2025 poster ()
Keywords: Causal Machine Learning
Personalized Survival Models
Time-to-Event Analysis
Treatment Effect Heterogeneity
Dose-Response Modeling
Primary area: machine_learning_for_sciences

Abstract: Estimating heterogeneous treatment effects (HTEs) of continuous-valued interventions on survival, that is, time-to-event (TTE) outcomes, is crucial in various fields, notably in clinical decision-making and in driving the advancement of next-generation clinical trials. However, while HTE estimation for continuous-valued (i.e., dosage-dependent) interventions and for TTE outcomes have been separately explored, their combined application remains largely overlooked in the machine learning literature. We propose DoseSurv, a varying-coefficient network designed to estimate HTEs for different dosage-dependent and non-dosage treatment options from TTE data. DoseSurv uses radial basis functions to model continuity in dose-response relationships and learns balanced representations to address covariate shifts arising in HTE estimation from observational TTE data. We present experiments across various treatment scenarios on both simulated and real-world data, demonstrating DoseSurv's superior performance over existing baseline models.

DoseSurv: Predicting Personalized Survival Outcomes under Continuous- Valued Treatments
Moritz Gögl
Yu Liu'
Christopher Yau!
Peter Watkinson'
Tingting Zhu
University of Oxford
Heelth Data Research UK
moritz. gogl @keble. ac.
Abstract
Estimating heterogeneous treatment effects (HTEs) of continuous-valiee inter- ventions on survival, that is, time-to-event (TTE) outcomes, is crucial in various fields, notably in clinical decision-making and in driving the advancement of next- generation clinical trials. However, while HTE estimation for continuous-vallue (i.e., dosage-dependent) interventions and for TTE outcomes have been separately explored, their combined application remains largely overlooked in the machine learning literature. We propose DoseSurv, varying-coefficien network designed to estimate HTEs for different dosage-dependent and non-dosage treatment options from TTE data. DoseSurv uses radial basis functions model continuity in dose- response relationships and learns balanced representations to address covariate shifts arising HTE estimation from observational TTE data. We present experi- ments across various treatment scenarios on both simulated and real-world data, demonstrating DoseSurv's superior performance over existing baseline models.
1 Introduction
Estimating the causal effects of interventions at the individual level is essential for informed decision making in fields such as personalized medieine [48], public health [39], social and economic policy [62], and marketing and retail [44]. Within the machine learning (ML) literature, predominant attention has been directed toward estimating heterogeneous treatment effects (HTEs) on continuous and binary outcomes [1, 2, 62, 70, 63, 60, 9, 16]. By contrast, time-to-event (TTE) outcomes, which capture the time until an event of interest, have received comparatively less attention. More specifically, the existing ML literature on HTE estimation with TTE outcomes has focused on quantifying causal effects of binary (non-dosage) interventions with "treatment" "no treatment" (control) being the only options [14, 58). However, in practice, treatment settings are typically more complex, comprising several, possibly continuous-valued treatment options, characterized by their dosage frequency of administration.
TTE outcomes have immense practical relevance, for example, in clinical studies that analyze survival time or time to discharge. Equally, studying the causal effects of continuous-valud treatments is an important problem healthcare, with relevant areas including chemotherapy [18], insulin regimens [36], and cardiovascular treatments [21]. Underlying dose-response relationships can be complex and monooonoo,ii, with multiple stimulatory and inhibitory phases. Although average respones curves across a population can be characterized in randomized controlled trials (RCTs) and may appear simple, dose-nesponse relationships at the patient level are often more intricate. Individual patient characteristics can complicate these relationships, leading to patient-specing ideal dosages and necessitating continuous modeling of treatments at the individual level.
While neural network (NN) models for HTE estimation have been developed separately for continuous- valued interventions [60, 9, 51] and for TTE outcomes [12, 14, 58], their joint application remains
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

largely unexplored in the current ML literature. Therefore, we propose DoseSurv, a deep learning survival model that estimates the conditional hazard function-the probability of experiencing the event at given time-under different treatments and dosages. Unlike many (semi-)parametric survival models-such as the Cox proportional hazards model [13] and related treatment-effect models [5ur model does not require covariate effects on the hazard to be constant over time.
Instead, we adopt NN-based discrete-time survival models that estimate conditional hazards through a series of binary classification tasks [10, 24], following their adaptation for HTE estimation with binary treatments by Curth et al. [14]. From the estimated conditional hazards, we can derive more practically relevant measures, such as treatment-specine survival time, which in turn allow us to define treatment effects and dose-response relationships. The core component of DoseSurv is the radial basis function (RBF) hazard estimator, which mod- continuous function of dosage via Gaussian RBFs,
survival functions and the (restricted) mean
els the survival hazard as treatment-spect leveraging the varying-comm ceeiingcoefficient network framework introduced by Nie et al. [51]. Through shared representations, DoseSurv also draws on information across different treatment options. Furthermore, it incorporates balanced representation learning via integral probability metric (IPM) regularization to mitigate various covariate shifts encountered in TTE and observational datasets [62, 14]. Here, event occurrence (survival), censoring (e.g., loss to follow-up), as well treatment and dosage assignment, typically depend on individual patient characteristics leading to different covariate shifts [14]. We therefore evaluate DoseSurv real and synthetic datasets spanning these covariate shifts and different treatment scenarios to assess its ability to mitigate the resulting biases.
Contributions Our contributions are fivefold: (i) We address the problem of HTE estimation for single- and multi-arm treatment regimens with mutually exclusive continuous-valus (dosage- dependent) and non-dosage treatment options, as well as TTE outcomes-a a setting that has been largely overlooked despite its high practical relevance. (ii) We propose DoseSurv, a novel varying- coefficient network based on layer-specific RBFs and designed for HTE estimation in settings such as the one described above. (iii) We adopt a computationally efficient IPM regularization technique to mitigate various biases arising in HTE estimation from observational TTE data and continuous valued treatments. (iv) Moreover, we create synthetic TTE datasets for different continuous- valued treatment types. and (v) we propose extension of the classical Twins dataset, as a real-world benchmark for continuous-valle valued treatment settings with TTE data. Our results show that DoseSurv outperforms existing baselines under different treatment scenarios and covariate shifts.
2 Related Work
ML methods for estimating HTES from binary interventions are well-studied and include model- agnostic meta-learners (e.g.. S- and T-learners) [40, 47, 72] and model-specific approaches like multi-task architectures that leverage shared information across treatment groups while also allowing trsetfent-ppeciffe modeling [62. 16, 2]. Implementations typically span tree-based models [69, 6] and NNs [34, 62, 63, 16]. and most works consider binary interventions only.
For cadinuuuu-valledd interventions, classical approaches use the generalized propensity score [30], while recent NN-based architectures directly model rese-response relationships. Some methods discretize dosage into bins with treatment-specife heads [60], offering regional flexibility at the expense of strict continuity; others adopt saryynn--oefficien change smoothly with dosage [51, 61]. preserving continuity. Adversarial approaches have also been explored for continuous-valuus interventions [9]; in TTE settings, however, censoring and temporal dependence complicate defining faithful discrimination targets and stable training objectives.
formulations in which parameters
In survival analysis, traditional models such the Cox proportional hazards model [13] have been extended using NNs [20, 37J. -or-mmramtttic alternatives include random survival forests [32] and discrete-time NN approaches [10, 45, 24], alongside survival clustering methods that group samples with similar risk profiles [33]. Integrating HTE estimation directly into survival models remains comparatively underexplored: most existing work addresses binary treatments using tree- rreebbsed [29] or NN approaches [14, 58], with no support for continuous- valuued treatments.
We provide an extended review of methods for survival analysis and HTE estimation-inchh incluudin those targeting continuous- valued interventions-in Appendix D, and summarize the applicability of related neural approaches across settings in Table D.I.

---

3 Problem Statement
3.1 Definitions
We aim to analyze an observational TTE dataset D comprising N individuals, each of whom is treated with one of M different continuous-valuu treatments. Throughout, largely follow notation and discrete-time survival formulations similar to those in [14]. The dataset can then be represented D { F2 8; (ai, :n where each patient is realization of random variables (x, A, (A, Q)) P. In this context, X X' refers to a vector of random variables associated with the patient covariates. Furthermore, we use T E T and C' E T to represent the random variables for the event time (e.g., the time until death or onset of disease) and the censoring time (i.e., the time until loss of follow-up). The observed patient outcomes are given by A = 1(T I C²), indicating whether either the event or censoring was observed, and T mTn(TT C), representing the time of that observation. We denote the random variable for the treatment assignment as E A (o...... M-1]. Each of these treatments A E A is associated with continuous random variable Q E Q - [0. 1] C R, representing the (normalized) dosage for the respective treatment." We treat time as discrete, with a finite time horizon tm such that T {1, fm Based the notation above, the probability of the event occurring at time T for a given patient with covariates a and intervention (a, q) who has not experienced an event prior to is defined as [14]
X(T az, a, 9) P(T A = = q, r),
(I)
which is also known as conditional hazard. From this, we can compute the survival function S(r|ac, a, q) = [les- X(1 ær, a, g)), representing the probability that the event will not occur up until time 7. To infer the conditional hazards (and thus the survival function), we follow a typical approach in discrete-time survival analysis by estimating A(r|æ, a, q) from covariates and time intervals using non-parametric ML model. We transform the TTE data into person-perior (long) format [66] and treat the estimation of the discrete conditional hazard function as a series of binary fit NN
classification problems associated with different time intervals. More specifically, model to the binary event indicator at each discrete time interval, with æ, a, 9, and T as predictors.
3.2 Relationship to Potential Outcomes Framework
In HTE estimation, our aim is not primarily to estimate outcomes conditional on the intervention, as presented above. Instead, we want to estimate the potential outcome that will occur in the event that particular action is performed. Following the Rubin causal model [31], we can adopt the survival function as a potential outcome of interest under a potential intervention (a, q) [14]
(T)( = P(T > r x| = x, do(A = a, Q = q, C 2 T)),
(2)
where the do-operator [53] shows that a specific action is being taken. Here, we have to treat censoring as an action we can intervene on by setting C 2 so that potential outcomes reflect uncensored event times and causal effects are identifiable [64].
The HTE between two potential interventions (a, q) and (ā, ②) may then, for example, be defined as the (time-dependent) difference in survival probabilities Vsurv (r(æ) (T)T( - ( (x). Alternatively, we may compute the restricted mean survival time (RMST) until the time horizon im under an intervention (a, q) as time-independent potential outcome of interest: ba
Et- (t/wx) The HTE can then be defined as (m)s(() = a.qq ( (x) 18.9 (x)) [14]. Moreover, for fixed treatment a, the function q + &ag (x) defines patient-specific dose-response curve that summarizes expected survival as a function of dosage.
To identify and quantify HTEs and dose-responses from TTE data, we make the following as- sumptions, which are standard the literature on ddosage-deeendeen treatment effect estimation [9, 60, 14] (Assumptions 1 1-3) and TTE outcomes [19, 50, 14] (Assumptionss 4-5), and extend them for continuous-valuse treatments. We denote the potential event time under a potential intervention (a, 4) Ta.q.
For simplicity, we focus here on a set of continuous- treatments only, i..., AA = _ coar However, DoseSurv also accommodate the more general case with continuous valued (dosage-dependee ical (nonedo treatment options, i.e.. A = A for the latter. q is a dummy with no effect.
and categor-
In the following, we refer to the tuple (a, g) as "intervention" to distinguish it from a "treatment" with associated "dosage"
3

---

Assumption 1 (Uneonfoundeedness) For all E A and E Q. the intervention and the potential event time are conditionally independent given covariates: (A, Q) Ta.g X.
Assumption 2 (Positivity of Interventions) For all X' and all e A: (2a) Positivity of Treatment: There exists > 0 such that P(A X m) <
(2b) Positivity of Dosage: For any admissible dose q E 2 and any small interval C Q around g. P(Q E u A a.
X = x) 0.
Assumption 3 (Consistency) The potential event time will actually be observed under the observed intervention without being affected by any other external factors, i.e.,
and Q = g. then
T = Ta.q
Assumption 4 (Independent Censoring) For all E A. Q, the potential event time and the cen- soring time are conditionally independent given covariates and intervention: Tag 1 X, (A,Q)
Assumption 5 (Positivity of Events and Censoring) For all æ E X, E A. q E Q, and all & E T:
(5a) Positivity of Events: P(T t X x. A = = q) 2 for some > 0.
(5b) Positivity of Censoring: IPC' > t X = x, A = a, Q = q) 2 c for some é > 0.
Under Assumptions 1, 3, and 4, we can identify the interventional hazard from the observational
T 2 r, X = æ. d = Q = 9, C,C. T)). The corresponding potential outcome survival function can then be computed as Sa.q (T
one, i.e., A(T|ar, 0, q) \T.z(T(az) IP(T
IIrsr (1 Aa.g (f)=2))))) (f)=)) A brief mathematical derivation of this identification is given in Appendix Assumptions and 5 further provide the overlap and non-degeneracy conditions for treatment, dosage, and at-risk sets required for non-parane parametric estimation [14]. We provide the causal diagram underlying the assumed data- gentraennn process in Appendix
DoseSurv
4.1
Model Architecture
Based on the definition of the problem presented above, we propose an NN-based discrete-time survival model, called DoseSurv, which is depicted in Fig. 1, and is tailored for survival predictions in single- and multi-arm treatment regimens with dosage-dependent (and non-dosage) treatments. DoseSurv consists of representation network, & X' + R, which is implemented as fully connected NN and shared across all treatments a, dosages g. and times т. It is followed by treatment- specific hazard estimators, implemented as individual networks, each performing multi-ouiput binary classification h(()) R x Q -> [0, fl which allow us to estimate the survival bazards at each time valuud
point T for a specific continuous- valued treatment a
Å(T|æ, a, g) = A/ (f(x),
where 115") is the output of (1(w) corresponding to time r. The treatment-specific estimators enable the learning of features unique to each respective treatment option. To enable the learning of the dosage dependency in the hazard estimates, we take two steps: (1) we add as an additional feature to the shared representation (z). (2) more importantly, we use Gaussian RBFs to explicitly model the network parameters (weights and biases) as continuous functions of as described below."
4.2
The RBF Hazard Estimator
For iattinuou-vvalued treatments a, want preserve the prominent role of q and prevent its influence from vanishing in a high-dimensional feature space. Therefore, we implement the traammen-speecfiic hazard estimators flexible carying-coefficient networks, where the network parameters are functions of dosage following prior work by Nie et al. [51]. We depart from that formulation by modeling the network parameters as linear combinations of treatment- and layer-specific RBFs with learnable centers and bandwidths. Concretely, let the hazard estimator
For non-dosage treatments E A ( (1) : - 10, 1] implemented standard fully connected NN.

---

Repreontation Netwaek
REP thaaadd Extimator
v(()
risk
(x
Stratification (1
(00C)) spent (oie
Survival Curves
Dose Besponse Curves
6%
920
LIPM
Tinet
Desage
Figure 1: DoseSurv features shared representation network that captures common features across treatments, dosages, and time intervals. To address covariate shifts arising from treatment-sclection bias, dosage-selection bias, censoring. and event-induced shifts, we introduce an IPM regularization term based on a custom discretization of dosage (4 q) and time (f + t), enabling the learning of balanced representations across treatments, dosages, and event times. For continuous treatments, A the model extends to the RBF hazard estimator, which consists of -reatient-speciffc network heads. These heads use treatment- and layer-specific RBF bases { J to model parameters as continuous functions of q, ensuring continuity of discrete-time hazard estimates with respect to dosage. The hazard estimates allow for the computation of treatment- and dosage-specific survival curves, enabling dose-response analysis.
((()) comprise Ji parameters (weights and biases) in each layer E (1,..., LJ. which are given by e (q) (0(.) (q). ())(()) (q)...., ( (q)) Each these parameters, indexed by is modeled as
continuous function ofq
K
0
p(a) (c) (9) Σ
(4)
k=1
{a) K where { 1==1 are the K basis functions in layer / for treatment a, and 8/1 are the learnable
coefficients [51]. We can therefore write
(g) B.(e) 1p (g)
(5)
with
((a)
E RAK
(1() = (win) (4), (1) (.....)))))))))))))) ERK
gin) Jatk
We propose the use Gaussian RBFs as basis functions
(()) ((()) (g) = (0) exp [ (q - -a- (2(0)))
(6)
with learnable centers (c
9, .... (.K) shape parameters (0,1 1.2
(a)
(a) (a)
........0 o.k) and additional
(a)
weights (w((), w/?)...,, w ). This flexibility, together with the universal approximation property of Gaussian RBF expansions 4, implies that-given sufficiently many centers and suitable band- widths-our g-dependent parameterization can represent any continuous dose-response function on compact intervals [52]. Moreover, Gaussian RBFs provide global support
the dosage range, potentially aiding in low-overlap regions. Furthermore, because complex survival likelihoods and their gradients interact nonlinearly with basis functions, our approach aims boost optimization by learning not only the varying coefficients but also the properties of the basis functions themselves.
'While not explored here, coupling these amplitude weights across clinically similar treatments would provide simple route induce soft cross treatment information sharing between hazard estimators.
5

---

4.3 Loss Function and Representation Balancing
DoseSurv uses both empirical risk minimization and representation balancing to learn hazard estimates in the presence of biases in treatment and dosage selection, as well as censoring and event occurrence.
Empirical Risk Loss
We minimize the empirical risk through the following loss term [24, 41]
risk
1
Σ BCE (b:(t), h( (+)))).
N
tm t=1 f::22
where (1( = 1 16; = T = =} denotes the observed event indicator for patient i in person- period (Iong) format. In this formulation, Liss is proportional to the negative discrete-time hazard bog-lielihooo-thar is, the average binary cross-entropy (BCE) computed over at-risk rows only. Minimizing it corresponds maximun-likelihodd estimation of per-interval hazards, and censoring is handled via the at-risk conditioning 2t
Representation Balancing To mitigate the effects of confounding, well censoring and event- induced shifts, DoseSurv aims to learn a representation & (x) in which features are balanced across at-risk populations at different times under different treatments and dosages. For this purpose, adapt an IPM regularization scheme for binary treatments [14], to minimize the Wasserstein distance [68] between feature distributions of individual time-, treatment-, and dosage-specific subpopulations and the overall baseline population. We discretize the continuous dosage into Q equally-sized bins {0,, 1/2), [1/2, 2/2). [(22 1)/2, I]]. indexed by q E {1, ... .22, ). and denote by q the dosage-bin index for sample i. To prevent oversegmentation into overly small subpopulations-whiin can reduce computational efficiency, generalizability, and stability- we also discretize time / into Т intervals of lower temporal resolution, indexed by E 21, .... I). Let & be the coarse time-bin index for sample i, obtained by mapping the observed time ñ to one of these I intervals. We then define the IPM regularization loss as the sum of Wasserstein distances between the representations of the total baseline population and that of each subpopulation
1 111 ΣE
LIPM
Wass( ((sss((((::)))) (m
A GEAq=1
Consequently, LIPM penalizes, for each treatment a, time bin t, and dosage bin q, the discrepancy in & between the corresponding at-risk subpopulation and the overall baseline population." In principle, one could complement or replace our IPM regularization with observation reweighting tailored dynamic survival settings with continuous treatments. Such an approach would require estimating generalized propensity scores for treatment and dosage given covariates, as well as conditional at-risk (survival) probabilities over time. In practice, these quantities are difficult to estimate
with continuous dosages and time-varying risk sets sets-and often yield extreme or unstable weights under limited overlap, which can inflate variance and destabilize training. To avoid these issues, we focus on IPM-based representation balancing.
We calculate the total loss function the sum of the empirical loss and the IPM regularization term
C = Lria -IPM-
(9)
Here, 7 is hyperparameter that controls the influence of the regularization term. More details on the implementation of DoseSurv provided in Appendix E.
5
Experiments
5.1 Experimental Setup
Following previous work on HTE estimation with continuous-valued treatments [60, 9] and TTE outcomes [14, 58], we first evaluate the performance of our model on synthetic data, which provides access to counterfactual outcomes and ground-truth survival curves for individual patients, and is therefore essential for a thorough evaluation of causal effect estimation. In addition, we evaluate DoseSurv real-world data, namely the Tivins dataset.
For non dosage treatments At stratification of at-risk populations reduces to (o, t): we retain the same normalization keep IPM loss contributions comparable across cantinuous-valuee and non-dosage arms.

---

5.1.1 Synthetic Datasets
We consider three simulated treatment scenarios (S1-S3) and create synthetic TTE datasets with different underlying outcome generation processes for each of them. Patient covariates are correlated in each of the datasets. In S1, we evaluate model performance for single continuous-valued treatment, i.e., Asi {0}. In S2, we consider two continuous-valle
treatment options Ase = {0. }}. Finally, we investigate DoseSurv's s performance in a setting with one non-dosage treatment option AS3 (0} and one continuous- valued treatment option Ass
Agoat {1). Extensive descriptions of the
data generation process provided in Appendix A.I. Our dosage assignment mechanism includes a parameter 7 that controls the level of dosage selection bias in the observed policy. When 1) = 0, the dosage is sampled from uniform distribution between [0, 1]. As 1) increases, dosage selection becomes increasingly confounded. Throughout all experiments, we simulate dosage assignment for ? (0, 2, 4). In addition, simulate fixed treatment selection bias for scenarios S2 and S3, which comprise multiple treatment options, and employ an informed censoring process, which introduces an additional censoring bias over time.
5.1.2 Ablations on DoseSurv Design
To examine the impact of certain design elements of DoseSurv, we conduct ablation studies. We compare the standard DoseSurv model (5 RBFs with layer-specifes learnable centers initialized at {0, 0.25, 0.5, 0.75, 11) with different variants: The first ablation (A1) corresponds to the standard DoseSurv model. except with fixed (non-Icarnable) RBF centers. For the second variant (A2), replace the RBF basis with a truncated polynomial (power) basis
for k ......d 0)d for k d 1, 1....,d.. d+
VIA (
max(q C
(a)
(10)
to construct splines of degree d = 2. with = 5 learnable knots, c. CL initialized at {0, 0.25, 0.5, 0.75, 0.95). This results in total of K basis functions, and thus higher plexity than the standard DoseSury model. The third ablation (A3) corresponds to the standard DoseSurv model, but with only three basis functions centers initialized at {0.2, 0.5, 0.83). For ablation A4, we do not add q as an additional feature to the shared representation (x)). incorporating only implicitly by modeling the network parameters using the RBFs that depend on q. In the final ablation (A5), we employ the
network design proposed by Nie et al. [51] for non- non-TTE data. This design features a truncated polynomial basis of degree = 2 with x = non-learnable knots at { {1/3, 2/3], and omits as an additional input feature. This results in K =5 basis functions, yielding complexity comparable to our standard DoseSurv model.
5.1.3 Real-world Dataset
Additionally, we evaluate DoseSury on the Twins dataset, a well-established benchmark in the HTE literature for binary treatments. This dataset provides survival times for twins born in the United States between 1989 and 1991 [3]. The original use of this dataset focuses on binary treatment "esignment-classffyy
the heavier twin as "treated" and the lighter twin "control" [46, 70J. First, we keep this notion and validate DoseSurv in a binary treatment scenario, without dosages, similar to [14]. Second, we extend the scope of the dataset by replacing the binary treatment indicator with the actual birth weight, treating birth weight as a continuous "dosage". This reinterpretation allows the dataset to represent two potential outcomes for each sample exposed to two distinct dosages. Details on the dataset and the implemented censoring mechanism are provided in Appendix A.2.
5.1.4 Baselines
We benchmark DoseSurv against well-established and sfateoof-the-ar ML models for TTE analysis:
DeepSury [37]: A continuous-timg survival model that extends the Cox proportional hazards model by incorporating NNs while retaining the assumption of proportional hazards.
DeepHit [45]: A discrete-time NN method modeling the event-time distribution, trained with composite likelihood and ranking losses; permits non non-propottional hazards.
RSF (Random Survival Forest) [32]: non-parametrie MI. ensemble of decision trees (bootstrap sampling, log-rank splits) for right-censored survival data.
7

---

NSC (Neural Survival Clustering) [33]: A method combining NNs and clustering to model TTE data by grouping samples with similar survival patterns and learning hazard functions per cluster.
By default, we implement all baselines single models (S-learner) with treatment and dosage information as additional covariates. For scenarios S2 and S3 where multiple treatment options are available, we also compare against treatment-specifee versions (T-learner) of the baseline models described above-traines separately per arm with the corresponding dosage included as an input feature-and and refer them DeepSurv-T, DeepHit-T. RSF-T and NSC-T. Moreover, we consider a version DoseSurv without IPM regularization, DoseSurv (no IPM). For the binary treatment scenario the Twins dataset, additionally compare DoseSurv against existing binary HTE survival models, namely SurvITE [14] and BITES [58].
5.1.5 Performance Measures
For simulated data, have access to the ground-truth individual dose-response relationships for every treatment a and dosage q E [0, 1]. We can therefore quantify how closely the estimated patient-specific dose-response curves match the true underlying curves across dosages and treatments using the mean integrated squared error (MISE) [60, 9]. with RMST as the outcome of interest
N
MISE
ΣΣ REAi=1
(Pap (2;) - taan (2;))) dg.
(11)
|A] N
For continuous-valued treatments a E A we also evaluate the mean dosage policy error (DPE) [60]. i.e., the mean squared error between the ground-truth outcome under the estimated optimal dosage ai arg maX frust (æ;) and the true optimal dosage Ff arg max euma (æ.). maximizing RMST. This is given by
DPE
1 N
Σ
(rrmst
(x:;)
R, (a::))
(12)
Accon
BEAA
The performance measures described above require knowledge of the outcome generation mechanism and are thus impractical for real-world applications. Therefore, we also provide common metrics in survival analysis based observed (factual) outcomes: the time-dependent C-index (Cd), the integrated Brier score (IBS), and the integrated negative binomial log-likelihood (INBLL). Details on these metrics and corresponding model performance are provided in Appendix B.
In the Twins dataset adapted for continuous- valued treatments, for each sample i, we have access to the event times T) -To.y.o and 7( -To.g under the lower dosage (birth weight) g and the higher dosage q. (2) of the same treatment AA {0], but not across the whole dosage range. Instead of computing MISE, we therefore provide the MSE between the observed and the estimated HTE
N (Dima,) Vrmst,i)²
MSE
N Σ
(13)
where V'ma,; min (TJ.g (m)) min (()) ( () (.m) and 0imaa,s = -ome ( (xx) are the observed and estimated HTEs in terms of the RMST. Analogously, we compute the MSE between estimated
frmus
0.9!
and observed HTES for the non-dosage (binary) treatment case, A = A {0, 1}. where 7: represents the survival time of the heavier twin (treated) and T)(() of the lighter twin (control).
5.2 Results and Discussion
5.2.1
Main Results
Fig. 2 shows model performance across five runs for scenarios S1-S3 under varying levels dosage selection bias (n (0, 2, 3, 4)). In all experiments, either DoseSurv or its derivative without IPM regularization consistently achieved the best performance in terms of both MISE and DPE,
Forr non-dosaa dosage treatments a A emma (x) is constant over reducing MISE to regular MSE between predicted and ground-truth RMST.
8

---

770
10-
$3.0
3.
DoseSurn
DeepSury
NSC
DoseSurn PM)
DeepSuruT
DeeaHi-T
ASPT
NSCT
Figure 2: Model performance for scenarios S1-S3 under different levels of dosage-assignment bias, averaged across 5 runs. Smaller values for v MISE and vDPF indicate better performance.
outperforming all baseline models. As expected, performance generally declined slightly with increasing for both metrics. Notably, IPM regularization substantially improved performance in most scenarios, particularly in S2, which involves two continuous- treatments, and primarily in terms of MISE. IPM regularization also yielded performance improvements even in the absence of dosage-selection bias (1) 0) and treatment-selection bias (S1), We attribute this to its ability to mitigate other covariate shifts inherent in TTE data, such event- or censoring-induced shifts.
The baseline models did not exhibit a clear performance trend. Among them, RSF and RSF-T generally achieved the lowest DPE. NSC performed comparably well in terms of DPE for S1, but showed the worst MISE performance in S2 and S3. DeepSurv(-T), DeepHit(-T), and NSC-T consistently demonstrated similar performance across all metries and experiments.
Table summarizes the performance on the real-world Twins dataset. Modeling the problem as a continuous treatment scenario and incorporating dosage information substantially improved performance across models, highlighting the importance continuous treatment modeling in survival analysis. DoseSurv demonstrated the best performance in estimating the HTE in both binary and continuous treatment settings, with substantial performance improvement observed continuous Among the baselines, DeepHit achieved the highest performance in the continuous setting, while RSF and DeepHit performed equally well in the binary setting.
Table Model performance on the Twins dataset across 5 runs (mean-195% CI). For baselines, we also show the relative difference compared DoseSurv.
Twins (continuous)
Twins (binary)
Methods
VMSE
VMSE
DeepSurv
14.461::0.217 (+15.7%)
16.235100111 (+4.1%)
DeepHit
13.272++0.200 (+6.2%)
15.97410.090 (+2.4%)
RSF
14.230+t0.092
(+13.8%)
15.923+-00022
(42.1%)
NSC
14.288+00555
(+14.3%)
16.694+-01223
(47.0%
BITES
16.177--0.161
(43.7%)
SurvITE
19.59914.725
(+25.6%)
DoseSurv
12.501-10.063
15.60110.531
5.2.2 Ablation Study Results
Figure illustrates DoseSurv's performance for scenario S2 (17 1), compared various ablations (A1-A5), described in Section 5. 1.2. The proposed design consistently achieved the best performance across all metrics compared other variants. In particular, each individual design choice yielded performance improvement over a DoseSurv implementation based on traditional network design. Moreover, the of RBFs results in fewer basis functions and coefficients, g(a) and thus lower model complexity, for the same number of centers/knots.

---

KGG
0.6
AT Leervooe
A2
A3 61092 O 68)
A4
A5 Coe Confirent Nowwork
Bosodurv
Figure 3: Impact of design choices (AI-A5) on the performance of DoseSurv in scenario S2 (7)=1 1) averaged across 5 runs. All other hyperparameters were held constant. Performance metrics were normalized such that the highest mean value across all configurations equals 1. Lower values imply better performance. We indicate model complexity in terms of the number of basis functions, K.
6 Conclusion
In this paper, we presented RBF-based varying-coefficient network for estimating HTES from TTE data in single- and multi-arm treatment regimens with continuous-valuee (dsnegddeeee)))) and non-dosage treatment options. We demonstrated in various experiments the benefit of the design in TTE settings with continuous treatments. Furthermore, we demonstrated significant performance improvements by incorporating IPM regularization scheme tailored for these settings, effectively mitigating covariate shifts caused by confounding factors and the nature of TTE data.
Limitations Like all HTE methods, model relies certain assumptions. In practice, leveraging expert and domain knowledge helps assess these and understand the impact possible violations. Another challenge tuning hyperparameters- expeccilly the IPM regularization strength y-under covariate shift. One way to address this is to select the largest y that preserves discrimination and calibration (e-g., via INBLL) on the validation Finally, as with most deep learning models, performance benefits from sufficient training data; our sample size ablation in Appendix C.I indicates degradation with small training sets.
balancing regularization and predictive power [14].
Broader Impact promising future application of models like those explored in this work is adaptive clinical trials, where integrating observational TTE data and leveraging ML can optimize trial designs. This approach could enable dynamic adjustments in patient recruitment, treatment, and dosage options, ultimately improving trial outcomes and patient care. However, analyses must contend with covariate shifts-includng treatment- and dosagenselection bias in observational settings and event- or censoring- induced shifts inherent to survival data-which, if unaddressed, can bias effect estimates and hinder optimal care.
Acknowledgments
MG is supported by the EPSRC Centre for Doctoral Training in Health Data Science (EP/S02428X//). CY is supported by a UKRI Turing AI Acceleration Fellowship (EP/VO23233/1) and received additional funding through EPSRC grant EP/Y018192/1. TZ is supported by the Royal Academy of Engineering under the Research Fellowship scheme.
References
[1] A. M. Alaa and M. van der Schaar. Bayesian inference of individualized treatment effects using multi-task gaussian processes. In Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. dropout, 2017.
[2] A. M. Alaa, M. Weisz, and M. van der Schaar. Deep counterfactual networks with propensity-
10

---

[3] D. Almond, K. Y. Chay, and D. S. Lee. The costs of low birth weight. The Quarterly Journal of Economics, 120 Aug. 2005.
[4] L. Antoliii, P. Boracchi, and E. Biganzoli. A time-dependent discrimination index for survival data. Statistics in Medicine, 24(24))3927-3944 Nov. 2005.
[5] Assaad, Zeng, C. Tao, S. Datta, N. Mehta, R. Henao, Li, and L.. Carin Duke. Counterfac- tual representation learning with balancing weights. In A. Banerjee and K. Fukumizu, editors, Proceedings of The 24th International Conference on Artificial Intelligence and Statisties, volume 130 of Proceedings of Machine Learning Research, pages 1972-1980, PMLR, 13-15 Apr 2021.
[6] S. Athey, J. Tibshirani, and S. Wager. Generalized random forests. The Annals of Statistics, 47(2), Apr. 2019.
[7] E. Bair and R. Tibshirani. Semi-supervised methods to predict patient survival from gene expression data. PLoS Biol., 2(4):E108, Apr. 2004.
[8] A. Bellot, A. Dhir, and G. Prando. Generalization bounds and algorithms for estimating conditional average treatment effect of dosage, 2022.
[9] I. Bica, J. Jordon, and M. van der Schaar. Estimating the effects of continuous-valued inter- ventions using generative adversarial networks. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 16434-16445. Curran Associates, Inc., 2020.
[10] E. Biganzoli, Boracchi, L. Mariani, and E. Marubini. Feed forward neural networks for the analysis of censored survival data: partial logistic regression approach. Statistics in Medicine, May 1998.
17(10):1199-1186
[11] C. Bockel genannt Rickermann, T. Vanderschueren, Berrevoets, T. Verdonck, and W. Verbeke. Using representation balancing to learn conditional-average dose responses from clustered data, 2025.
[12] P. Chapfuwa, S. Assaad, S. Zeng. M. Pencina, L. Carin, and R. Henao. Enabling counterfactual survival analysis with balanced representations. In Proceedings of the Conference on Health, Inference, and Learning, ACM CHIL '21. ACM, Apr. 2021.
[13] D. R. Cox. Regression models and life-tables. Journal of the Royal Statistical Society Series B: Statistical Methodology, 34(2): 102-202, Jan. 1972.
[14] A. Curth, C. Lee, and M. der Schaar. Survite: Learning heterogeneous treatment effects from time-to-event data. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 26740-26753. Curran Associates, Inc., 2021.
[15] A. Curth and M. van der Schaar. Nonparametric estimation of heterogeneous treatment effects: From theory to learning algorithms. In A. Banerjee and K. Fukumizu, editors, Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Reseurch, pages 1810-1818. PMLR, 13-15 Apr 2021.
[16] A. Curth and M. van der Schaar. On inductive biases for heterogeneous treatment effect estimation. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 15883-15894, Curran Associates, Inc., 2021.
[17] M. Cuturi and A. Doucet. Fast computation of wasserstein barycenters. In E. P. Xing and T. Jebara, editors, Proceedings of the 31st International Conference on Machine Learning, volume 32 of Proceedings of Machine Learning Research, pages 685-693, Bejing, China, 22-24 Jun 2014. PMLR.
[18) L. R. Duska, R. Penson, J. G. Supko, D. M. Finkelstein, Makastorsis, J. Gallagher, K. Borden, A. Goodman, A. F. Fuller, N. Nikrui, and M. V. Seiden. A phase i study of continuous infusion doxorubicin and paclitaxel chemotherapy with granulocyte colony-stimulating factor for relapsed epithelial ovarian cancerl. Clinical Cancer Research, 5(6): 1209-130,, 06 1999.

---

[19] I. Díaz, O. Savenkov, and K. Ballman, Targeted learning ensembles for optimal individualized treatment rules with time-to-event outcomes. Biometrika, 105(3):723-738, May 2018.
[20] D. Faraggi and R. Simon. A neural network model for survival data. Statistics in Medicine, 14(1):73-2,, Jan. 1995.
[21] M. Fiuzat, D. Wojdyla, D. Kitzman, J. Fleg, S. J. Keteyian, W. E. Kraus, I. L. Piña, D. Whellan, and C. M. O'Connor. Relationship of beta-blocker dose with outcomes in ambulatory heart failure patients with systolic dysfunction. Journal of the American College of Cardiology, 60(3):208-211 200-211,,,a0 July 2012.
[22] J. A. Foekens, H. A. Peters, M. P. Look, H. Portengen, M. Schmitt, M. D. Kramer, N. Brünner, F. Jänicke, M. E. M.-v. Gelder, S. C. Henzen-Logmans, W. L. J. van Putten, and J. G. M. Klijn. The urokinase system of plasminogen activation and prognosis in 2780 breast cancer patientsl. Cancer Research, 6063):636-643, 02 2000.
[23] C. Foster, J. M. Taylor, and S. J. Ruberg, Subgroup identification from randomized clinical trial data. Statistics in Medicine, 30(24):8677-2880 Aug. 2011.
[24] M. F. Gensheimer and B. Narasimhan. A scalable discrete-time survival model for neural networks. Peerl, 7:c6257, Jan. 2019.
[25] T. A. Gerds and M. Schumacher. Consistent estimation of the expected brier score in general survival models with right-censored event times. Biometrical Journal, 4866):002--1440, Dec. 2006.
[26] E. Graf, C. Schmoor, W. Sauerbrei, and M. Schumacher. Assessment and comparison of prognostic classification schemes for survival data. Statistics in Medicine, 18(17-18):2529225 2545, Sept. 1999.
[27] N. Hassanpour and R. Greiner. Counterfactual regression with importance sampling weights. In Proceedings of the Twenty-Eighth International Joint Conference on Artificial Intelligence, IJCAI-19, pages 5880-5887. International Joint Conferences on Artificial Intelligence Organi- zation, 2019.
[28] Hastie and R. Tibshirani. Varying-coefffcient models. Journal of the Royal Statistical Society. Series B (Methodological), 55(4):757-796, 1993.
[29] N. C. Henderson, T. A. Louis, G. L. Rosner, and R. Varadhan. Individualized treatment effects with censored data via fully nonparametrie bayesian accelerated failure time models. Biostatistics, 21(1):50-6) 68, July 2018.
[30] G. W. Imbens. The role of the propensity score in estimating dose-response functions. Biometrika, 876))706-710, 2000.
[31] G. W. Imbens and D. B. Rubin. Rubin Causal Model, page 229-241. Palgrave Macmillan UK 2010.
[32] H. Ishwaran, U. B. Kogalur, E. H. Blackstone, and M. S. Lauer. Random survival forests. The Annals of Applied Statistics, 2(3), Sept. 2008
[33] V. Jeanselme, B. Tom, and J. Barrett. Neural survival clustering: Non-parametric mixture of neural networks for survival clustering. In G. Flores, G. H. Chen, T. Pollard, J. C. Ho, and T. Naumann, editors, Proceedings of the Conference on Health, Inference, and Learning, volume 174 of Proceedings of Machine Learning Research, pages 92-102. PMLR, 07-08 Apr 2022.
[34] F. Johansson, U. Shalit, and D. Sontag. Learning representations for counterfactual inference. In M. F. Balcan and K. Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 3020- 3029, New York, New York, USA, 20-22 Jun 2016. PMLR
[35] D. Kalbfleisch and R. I.. Prentice. The Statistical Analysis of Failure Time Data, Wiley, Aug. 2002.
12

---

[36] U. Kampmann, S. Knorr, J. Fuglsang, and P. Ovesen. Determinants of maternal insulin resistance during pregnancy: An updated overview. Journal of Diabetes Research, 2019:1-9, Nov. 2019.
[37] J. L.. Katzman, U. Shaham, A. Cloninger, J. Bates, T. Jiang, and Y. Kluger. Deepsurv: person alized treatment recommender system using cox proportional hazards deep neural network. BMC Medical Research Methodology, 18(1), Feb. 2018.
[38] D. Kingma and J. Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), San Diego, CA, USA, 2015.
[39] N. Kreif, K. DiazOrdaz, R. Moreno-Serra, A. Mirelman, T. Hidayat, and M. Suhrcke. Estimating heterogeneous policy impacts using causal machine learning: a case study of health insurance reform in indonesia. Health Services and Outcomes Research Methodology, 22(2):192-22)
Nov. 2021.
[40] S. R. Künzel, J. S. Sekhon, P. J. Bickel, and B. Yu. Metalearners for estimating heterogeneous treatment effects using machine learning. Proceedings of the National Academy of Sciences, 16610))445644165, Feb. 2019.
[41] H. Kvamme and Ø. Borgan. Continuous and discrete-time survival prediction with neural networks. Lifetime Data Analysis, 27(44:710-736, Oct. 2021.
[42] H. Kvamme and Ø. Borgan. The brier score under administrative censoring: Problems and solution. Journal of Machine Learning Research, 24(2):1-26, 2023.
[43] H. Kvamme, Ø. Borgan, and I. Scheel. Time-to-event prediction with neural networks and cox regression. Journal of Machine Learning Research, 20(129):1330 2019.
[44] H. Langen and M. Huber. How causal machine learning can leverage marketing strategies: Assessing and improving the performance of coupon campaign. PLOS ONE, 18())=e0278337 Jan. 2023.
[45] C. Lee, W. Zame, J. Yoon, and M. Van der Schaar. Deephit: deep learning approach to survival analysis with competing risks. Proceedings of the AAAI Conference on Artificial Intelligence, 32(1), Apr. 2018.
[46] C. Louizos, U. Shalit, J. M. Mooij, D. Sontag, R. Zemel, and M. Welling. Causal effect inference with deep catent-variabll models. In Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017
[47] M. Lu, S. Sadiq. D. J. Feaster, and H. Ishwaran. Estimating individual treatment effect in observational data using random forest methods. Journal of Computational and Graphical Statistics, 27(1) (1)2:22021199 Jan. 2018
[48] A. D. Meid, C. Ruff, I.. Wirbka, F. Stoll, H. M. Seidling, A. Groll, and W. E. Haefeli. Using the causal inference framework to support individualized drug treatment decisions based on observational healthcare data. Clinical Epidemiology, Volume 12: 12:1222-1234, Nov. 2020.
[49] M. Mesinovie, P. Watkinson, and T. Zhu. Dysurv: dynamic deep learning model for survival analysis with conditional variational inference. Journal of the American Medical Informatics Association, page ocae271, 11 2024.
[50] C. Nagpal. M. Goswami, K. Dufendach, and A. Dubrawski. Counterfactual phenotyping with censored time-to-events. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, KDD '22. ACM, Aug. 2022.
[51] I.. Nie, M. Ye, Q. Liu, and D. Nicolae. VCNet and functional targeted regularization for learning causal effects of continuous treatments. In International Conference on Learning Representations, 2021.
[52] Park and I. W. Sandberg. Universal approximation using basis-functiion networks. Neural Computation, 3(2):246-57,, June 1991.
13

---

[53] J. Pearl. Causality: Models, Reasoning, and Inference. Cambridge University Press, Sept. 2009.
[54] S. Pölsterl. scikit-survival: A library for time-to-even analysis built top of seikit-leam. Journal of Machine Learning Research, 21(212):166 6, 2020.
[55] R. Ranganath, A. Perotte, N. Elhadad, and D. Blei. Deep survival analysis. In F. Doshi-Velez, J. Fackler, D. Kale, B. Wallace, and J. Wiens, editors, Proceedings of the 1st Machine Learning for Healthcare Conference, volume 56 of Proceedings of Machine Learning Research, pages 101-114, Northeastern University, Boston, MA, USA, 18-19 Aug 2016. PMLR.
[56] K. Ren, J. Qin, L. Zheng, Z. Yang, W. Zhang, L. Qiu, and Y. Yu. Deep recurrent survival analysis. Proceedings of the AAAI Conference on Artificial Intelligence, 33(01):47998 4805, July 2019.
[57] P. Royston and D. G. Altman. External validation of cox prognostic model: principles and methods. BMC Medical Research Methodology, 13(1), Mar. 2013.
[58] S. Schrod, A. Schäfer, S. Solbrig, R. Lohmaaer, W. Gronwald, P. J. Oefner, T. Beißbarth, R. Spang, H. U. Zacharias, and M. Altenbuchinger. Bites: balanced individual treatment effect for survival data. Bioinformatics,
(attpplement_)x:i60---ie June 2022.
[59] M. Schumacher, G. Bastert, H. Bojar, K. Hübner, M. Olschewski, W. Sauerbrei, C. Schmoor, C. Beyerle, R. L. Neumann, and H. F. Rauschecker. Randomized 2 x 2 trial evaluating hormonal treatment and the duration of chemotherapy node-positive breast cancer patients. german breast cancer study group. Journal of Clinical Oncology, 12(10):2086-208
Oct. 1994.
[60] P. Schwab, L. Linhardt, S. Bauer, J. M. Buhmann, and W. Karlen. Learning counterfactual representations for estimating individual dose-response curves. In Proceedings of the 34th Conference Artificial Intelligence (AAAI), volume 34, pages 5612-5611 AAAI Press, Feb. 2020. AAAI Technical Track: Machine Learning.
[61] J. Schweisthal,, D. Frauen, V. Melnychuk, and S. Feuerriegel. Reliable off-policy learning for dosage combinations. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 67900-67924. Curran Associates, Inc., 2023.
[62] U. Shalit, F. D. Johansson, and D. Sontag. Estimating individual treatment effect: generalization bounds and algorithms. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th Inter- national Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 3076-3085. PMLR, 06-11 Aug 2017.
[63] C. Shi, D. Blei, and Veitch. Adapting neural networks for the estimation of treatment effects. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bun Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[64] O. M. Stitelman and M. J. van der Laan. Collaborative targeted maximum likelihood for time to event data. The International Journal of Biostatistics, 6(1) ()), Jan. 2010.
[65] O. M. Stitelman, C. W. Wester, V. De Gruttola, and M. J. van der Laan. Targeted maximum likelihood estimation of effect modification parameters in survival analysis. The International Journal of Biostatistics, 7(1):1-34, Jan. 2011.
[66] K. Suresh, C. Severn, and D. Ghosh. Survival prediction models: introduction to discrete-time modeling. BMC Medical Research Methodology, 22(1), July 2022.
[67] S. Tabib and D. Larocque, Non-parametric individual treatment effect estimation for survival data with random forests. Bioinformatics, 46(2):629-636) Aug. 2019.
[68] C. Villani. The Wasserstein distances, page 93-111. Springer Berlin Heidelberg, 2009.
[69] S. Wager and S. Athey. Estimation and inference of heterogeneous treatment effects using random forests. Journal of the American Statistical Association, 1 113(523): 122--1242, June 2018.
14

---

[70] J. Yoon, J. Jordon, and M. van der Schaar. GANITE: Estimation of individualized treatment effects using generative adversarial nets. International Conference on Learning Representa- lions, 2018.
[71] W. Zhang, T. D. Le, L. Liu, Z.-H. Zhou, and J. Li. Mining heterogeneous causal effects for personalized cancer treatment. Bioinformatics, 33(15))2372-2311 Mar. 2017.
[72] Y.-F Zhang, H. Zhang, Z. C. Lipton, L. E. Li, and E. P. Xing. Exploring transformer backbones for heterogeneous treatment effect estimation, 2022.
15

---

Appendix
A Data Description
A.1 Synthetic Data Generation
We generate covariates from correlated multivariate normal distribution, ~ No Σ) with
= (1 - -- p11" T and p = 0.2. We generate 20,000 samples with 40 covariates for scenario S1 and 20 covariates for S2 and S3. To simulate event times, randomly sample number of parameter vectors relevant to define the event-processe
each scenario, u 2 u ~ (0, 12) and normalize them 2 =u//llz I|. where || || is the Euclidean norm, and S E {S1, S2, S3}. In scenarios S2 and S3, involving two treatment options, we aim to simulate an event process that reflects shared characteristics across treatments (e.g.. common risk factors such age) while also incorporating treatment- -seennif properties in the event process. denote shared parameter vectors as of and treatment-spen specific parameter vectors as i Finally, for the three scenarios S1, S2, S3, we use the following event processes, where (.)) is the sigmoid function:
S1 (1 continuous-valueed treatment option)
(t)a()))) = 0. (02 x)) (0) a sin (2r
1))ttta 1)t/tm)-
3.5 sin 2/3) + og m²) sgn(ng) x²)²²)
$2 (2 continuous-valued treatment options)
0.1 if 0.50 ao sin(5xx+ 5200
(t,g(1/z)=
x) x) 2 52.0 nq mq²), 25
a.) i
(t/oat)
$2,1 sin + 5.22110 + 82,1 meq²) if:25 (ff a)) sgn(0
S3 (1 non-dosage and
contimuous-valued treatment option)
0.1)(ef²x), ift<5
0.20 $3.00 w 1), i
0.10(ef²x), ift<5 0.50 6 +255311 mg"). ift25
(t]r)
Fig. A. 1 (A) shows simulated ground-truth survival curves for scenario $1 for 6 different samples and dosages g = 0.2, 0.5, 0.8. Fig. A.1 (B) depicts the corresponding dose-response relationships, where define the outcome of interest the RMST. We define the covariate-dependent, informative, censoring process as
S (4)) = 0. (r)
where v is another randomly sampled and normalized parameter vector, introducing additional censoring bias, and therefore covariate shift over time. Furthermore, we assume administrative censoring fm 30. To assign dosages, follow approach similar [9] and [61]. More specifically, for each continuous treatment option, a, we sample the dosages from a Beta distribution
qa ~ Beta(īi, by;))
with byg = 1 + 7j. where the mode of the Beta distribution is given by the true optimal dosage, 4. maximizing the RMST for the respective patient under treatment a. This mechanism mimics realistic medical scenario in which the dosage assignment is informed by the true optimal
To maintain symmetry, sample da Beta(bag, with baz (ñ 1)/(1 gb) +2 7 if 0.5.
16

---

Dese-rasporee Curves
[145T
92
18
0.2
1.
Figure A.1: Simulated (continuous) ground-truth survival curves for scenario S1 for 6 different samples and dosages q = 0.2, 0.5, 0.8 (A). Corresponding dose- response relationships with RMST as the outcome of interest (B).
dosage. The parameter yj fj controls the dosage assignment, i.e., how much influence the covariates have the dosage assignment. With = 0, the dosage is sampled from a uniform distribution, while the dosage assignment becomes increasingly confounded towards the optimal dosage increases.
To simulate confounded treatment assignment in scenarios S2 and S3, we design mechanism that links the observed treatment type a with the individual patient covariates. For this, we sample M parameter vectors ua (0, 1 and normalize them v, u /uall Then, we standardize the covariates æ and compute the probabilities for observing treatment a
exp(xve,
Po
ZAEA exp( (vV))
Next, we compute the modulated probabilities
(pa)s
Da Σ ((a))))
where controls the treatment selection bias confounding, i.e., how much influence the covariates have on the treatment decision (with = indicating random treatment decision). For both scenarios S2 and S3 we choose { = 1.5 and sample the treatment from the multinomial distribution
a ~ Multinomial )oeA:
suitable for any number of treatment options. Figure A.2 illustrates the covariate shift in two selected covariates, I11 + 19, between the two different treatment groups in scenario S2. The histograms represent the at-risk populations different time points 0, 10, 20) under different values of
50
255
Figure A.2: Histograms of F11 + Fi9 for at-risk populations across time for different treatments,
= 0 and a = 1, in scenario S2, Covariate shifts arise with increasing time t and bias parameter g.
17

---

the treatment selection bias parameter (E = O, 1.5). As & increases, the covariate shifts become more pronounced. Additionally, shifts evolve over time due to differing event probabilities between treatment groups [14].
Across all scenarios, we employ 17,000 samples for training. 1,000 samples for validation, and 2,000 samples for testing.
A.2 Twins Data
The Twins data is sourced from all US births between 1989-1999 [3], where we extract data on twins only, similar to [46, 70, 14]. Moreover, we select twins weighing less than 1,100 grams, extract 45 features predominantly related to parental factors, pregnancy, and birth, and include only samples with no missing features. We analyze 30-day survival, retaining twin pairs with no death after day 30; twins without recorded death are treated as administratively censored at day 30. Moreover, we keep the original time resolution of day as the length of the time intervals (tm = 30). We create semi-synthetic observational TTE dataset by selecting one twin from each pair. For this, we use a similar mechanism as previously for the treatment assignment in the simulated scenarios S2 and S3: First, randomly sample two parameter vectors un u ~ _N( 12) and normalize them twin = =h /I u/ [uin I, vi u
twiin /| ||//"")] After standardizing the covariates æ, we define the
probability for observing the lower-weight weight twin or higher-weigee twin
exp(xve)
Pa (0.1) exp(xvelxin)
Next, we compute the modulated probabilities
(pa)s
No
(o.
(0.1) ((pa)))²
where we choose g = 20. We apply a synthetic censoring mechanism inspired by Schrod et al. [58] to introduce early censoring into samples which are not already administratively censored at day 30. As previously, we generate a normalized random vector uvens and compute a censoring score si for each patient by computing the dot product with the standardized covariates: gçens wiwin ar, After standardizing sf across all samples, we censor all remaining samples where st > 0 (i.e., 50% of the remaining samples) at a random fraction f; U (0, 1) of the sample's is true observed survival time. Note that following this process we simulate a scenario where the independent censoring assumption is violated.
The final cohort contains 5,601 samples. For our analysis, we use 70% of samples for training. From the remaining samples, we extract 30% for validation and 70% for testing.
B Metrics Based on Observed Outcomes
B.1 Definitions
In the following, we provide definitions of the additional metrics used to measure the performance of survival models based on observed (factual) outcomes. Here we treat time as continuous and interpolate accordingly.
Time-dependent C-index (C" Following the definition from [4, 43], we compute Old as
Cdd ((Š((F) (2;)) < S(, 2 řj. 1).
This metric is not limited to settings where the proportional hazard assumption holds, but can be biased since only uncensored individuals are taken into account [49].
Integrated (IPCW) Brier Score (IBS): The IBS [26. 25] can be interpreted as a measure of the mean accuracy of the predicted survival probability over time, adjusting for censoring by incorporating inverse probability of censoring weights (IPCW). is computed as
max(F;)
N $(11))) (1. (1- (1= > dl,
IBS-
max
((F.)
G(t)
18

---

where we denote the Kaplan-Mett Meierr estimator applied to the censoring times G(t) = P(C > t). Here, biases may arise due to dependencies between censoring distribution and covariates [42, 43].
Integrated (IPCW) Negative Binomial Log-likelihood (INBLL): The INBLL measures the performance of survival model in terms of discrimination and calibration, averaged over time. is calculated as [43]
max(f.)
N log[S(t
log(1 - S(1|2))) (+t & 1,,0=11
INBLL= max(f.)
+
dt.
N
G(t)
((F.)
We use this metric for hyperparameter optimization of the baseline models (see Section E.2).
B.2 Results Observable Metrics
Fig. B.1 illustrates the model performance, measured by C IBS, and INBLL, for scenarios S1-S3, averaged over five runs. The results align closely with the symthetic data metrics presented in the main paper (see Fig. 2). Overall, DoseSurv consistently outperforms other models across most metrics and scenarios, with its variant without IPM regularization typically ranking second. Among the baseline models, DeepSurv generally achieves the best performance in terms of Cd., IBS, and INBLL in scenario $1, while RSF-T demonstrates high performance in scenarios $2 and S3.
51
52
53
0.75
0.70
6.65
0.00
1.00
0.58
0.35
6.20
120
0.10
0.55
0.50
0.00
0.255
1.00
0.50
0,25
6.00
DeseSurv
DeepSure
Despltt Deepit
RSF
NSC NSC-T
DeseSurv
DeepSoret
Figure B.1: Model performance for scenarios $1-S3, in terms of Old, IBS and INBLL under different levels of dosage-assignment bias 1), averaged across runs. Greater values of Cdd and lower values of IBS and INBLL indicate better performance.
C
Additional Experiments
C.1 Ablation Study on Sample Size
We conducted ablation study examine the impact of training dataset size on model performance in scenario S1 (1) = 0). We performed experiments using 5 different training sample sizes: 2,000, 6,000, 10,000, 14,000 and 18,000. Each training dataset was independently and randomly sampled from the combined training and validation pool. The same hyperparameters as in the main experiment for were employed. For methods using early stopping (DoseSurv, DeepHit, DeepSurv, NSC), we reserved 10% of each sampled dataset for validation; RSF was trained on the full sampled set. The results are depicted in Fig. C. DoseSurv achieved the best performance across all sizes and metries, except at the smallest size.
19

---

0.7
03:00
0.20
6.3
0.53
0.15
2002
00:30
12005
14003
100.0
2032
00.00
LOCE0
14030
100.0
2030
6000
10683
14000
18025
230
60.0
10033
14.00
18330
Cempdory
852
Omadaa
Ooredar
Figure C.I: Model performance across 5 runs (mean CT) for scenario S1 (1) 0) for different training dataset sizes.
C.2 Binary Treatment Scenario: Hormone Therapy in Breast Cancer Patients
We outline a potential application DoseSurv to optimize binary treatment decisions, specifically hormone therapy for breast cancer patients. This external validation scenario has been previously explored by [57, 37, 58). Although binary treatment settings are not the main focus DoseSurv, the same approach could be applied to continuous-value's treatment problem.
Following the methodology established in the referenced literature, we use data from 1,545 breast cancer patients in the Rotterdam Tumor Bank [22]. From this observational dataset, we extract six features: age, menopausal status, number of cancerous lymph nodes, tumor size, and progesterone and estrogen receptor statuses. Since treatment assignment in the Rotterdam dataset is non-random, we use these data for model training (85% training, 15% validation).
For testing, we use data from 686 patients enrolled in randomized controlled trial conducted by the German Breast Cancer Study Group (GBSG) [59], ensuring external validation. The same six patient characteristics are extracted from this dataset. To simplify the analysis, survival times were discretized into 30 equally-sized intervals for DoseSurv and all baseline models.
The results, summarized in Table C. demonstrate that DoseSurv, although designed primarily for continuous-valle treatments, be applied effectively in binary treatment settings. In partieu- lar, DoseSury achieved better calibration accuracy (lower IBS and INBLL) than baseline models. Furthermore, IPM regularization led to improvements in model performance.
Table C.I: Model performance on the GBSG dataset across 5 runs (mean-195% CI).
Methods
Ctd
IBS
INBLL
DeepSurv DeepHit
0.683+0.004 0.680=000009
0.118:+0.004
0.384+0.012
0.121=00022
0.394=00055
RSF
0.685+10.001
0.122+10.011
0.403+000022
NSC
0.672-1.012
0.000
0.370-L0.058
BITES
0.005
0.115-10.004
0.381+0.008
SurvITE
0.665++0.037
0.187++0.072
0.539+L0.174
DoseSurv (no IPM)
0.670++0.008
O. 0.110-00.005
0.356++0.012
DoseSurv
0.688-0.002
0.098
0.338+0.0009
C.3
Experiments under Different Time Discretization
Like other discrete-time NN survival models, DoseSurv supports flexible discretization of survival times into any number of intervals. In our main experiments, we used 30 time intervals, aligning with common practice in the literature [45, 14]. However, DoseSurv can readily operate under different temporal resolutions. Here, we perform additional experiments on data simulated using finer time discretization of 60 intervals and compare DoseSurv to the continuous- time, CoxPH-inspired DeepSurv model, which may offer greater robustness. Table C.2 presents corresponding results for data simulated under Scenario SI (1) 1). We find that DoseSurv yields metrics comparable to those obtained with 30 intervals in the main experiment, whereas DeepSury continues to underperform relative to DoseSurv.
20

---

Table C.2: Model performance of DoseSurv and DeepSurv under Scenario S1 (1) with finer time discretization (60 intervals). For reference, we include performance metrics achieved by DoseSurv in the corresponding main experiment with 30 intervals. Metrics are averaged over runs.
Intervals
Method
MIIS
VDPE
Ctd
IBS
30
DoseSurv
4.46
7.97
0.75
0.16
60
DoseSurv
5,51
8.00
0.76
0.14
DeepSurv (CoxPH)
8.44
10.43
0.61
0.20
C.4 Sensitivity Analysis
To assess the robustness of DoseSurv to hyperparameter choices, we conducted a sensitivity analysis under Scenario S1 (n 1). We varied the learning rate (Ir), batch size, and IPM regularization strength (7), while keeping all other hyperparameters fixed. Results are averaged over 3 independent runs. The model demonstrates consistent performance across a broad range of settings, indicating relative stability with respect to these key hyperparameters.
Table C.3: Sensitivity of DoseSurv to learning rate, batch size, and IPM regularization under Scenario
S1 (1) 1). Metrics are averaged over 3 runs.
Ir
Batch Size
7
MISE
VDPE
Ctdd
IBS
0.001
500
0.01
4.58
7.94
0.75
0.16
0.001
500
0.10
4.48
7.88
0.75
0.16
0.001
1000
0.01
4.66
8.09
0.75
0.16
0.001
1000
0.10
4.70
8.07
0.75
0.16
0.005
500
0.01
5.47
8.39
0.73
0.17
0.005
500
0.10
5.73
8.43
0.72
0.17
0.005
1000
0.01
4.81
8.04
0.74
0.16
0.005
1000
0.10
4.95
8.10
0.74
0.16
C.5 Visualizing Dose-Response Curves
Fig. C.2 shows example dose-response curve predicted by DoseSurv in single under scenario S1 (1) = 0) for a selected sample red. The corresponding ground truth curve is depicted in green. In practice, dose-response curves can be obtained by evaluating the model for range of dosage values & and using the predicted RMST as the outcome of interest.
30
Ground truth
25
Prediction
RMST
20
15
10
0.0
0.2
0.4
0.6
0.8
1.0
Dosage
Figure C.2: Example of predicted and ground-truth dose-response curves under scenario S1 (1) 0).
21

---

D Extended Literature Review
D.I ML Methods for HTE Estimation.
Inferring HTEs of binary treatments has been studied extensively in the recent ML literature. Fol- lowing Curth and van der Schaar [16], methods can typically be categorized into model-specte and model-agnostic approaches. The latter, also referred to as "meta-learners" [40], can be implemented using any kind of ML method. The most common types of meta- meta-learners (indirect estimators) aim to obtain the HTE as the difference between regression estimates of the two potential outcome functions. T-learners [40, 47], fit two (or more) regression models (one for each treatment group) separately, while S-Iearners [23, 34, 40, 72] fit a single regression model where the treatment variable is added as an extra feature. Meta-learners have been implemented using wide range of ML models, particularly tree-based models [23,4] 47, 69, 40, 6] and NÑs [34, 72]. On the other hand, model-specific approaches rely on unique properties of the underlying ML model. These include approaches which could be described as hybrids between S- and T-learnes [15]. Such methods are typically implemented as "multi-task learning" models, which learn different treatment responses simultaneously, allowing some information flow between treatments [1, 2, 62, 63, 16]. Beyond architecture, balancing and reweighting strategies are also used to mitigate confounding and covariate shift, including IPM-basel representation balancing [62], counterfactual representation learning with balancing weights [5], pepeoraance-aaplinng reweighting [27]. and propensity-bes regularization [63].
HTE estimation of continuous-vallue interventions is less explored in the literature. Traditional statistical approaches typically rely the generalized propensity score [30]. The dose-response- networks proposed by Schwab et al. [60] extend the multi-task learning NN architecture for binary treatments, first presented by Shalit et al. [62], by subdividing the treatment-specific layers, into multiple network heads, each assigned to a different dosage interval. This facilitates the regression of the dose-response function across different dosage regions. A drawback of this architecture is that bin-specific heads produce separate outputs for different dosage intervals, breaking the continuity in the fose-response curves. Inspired by [28]. Nie et al. [51] aim overcome this limitation by network (VCNet) in the prediction head. Here, the network weights
employing a varying-coefficien are modeled as continuous functions of the dosage. Schweisthal et al. [61] extended this approach to treatments. A different approach by
learn the effect of combinations of multiple dosage-dependee Bica et al. [9] extends the GAN framework proposed by Yoon al. [70], and introduced a hierarchical discriminator specifically tailored for continuous valued treatment settings. However, leveraging adversarial learning for counterfactual survival analysis is challenging due to censoring and time structure, which complicate defining faithful discrimination targets and stable training objectives. Existing balancing methods for continuous-valued treatments assume clustered data [11], which limits its applicability in settings where natural clusters absent. Moreover, Bellot al. [8] provide a generalization bound, but their method relies on estimating complex integrals over representation and treatment spaces, which makes it likely to be sensitive to function and kernel choices.
D.2 Survival Analysis and HTE Inference with TTE Data.
ML methods for survival analysis have gained significant attention in recent years. Notable approaches include adaptations of the Cox- Cox-proportional hazards model [13], where traditional linear predictors are replaced with feedforward NNs [20, 37]. In addition, TTE has been modeled via accelerated failure time (AFT) regression [35] and deep generative models [55]. Random survival forests have non-parametric classifiers [10, 24, 56, 41] have gained popularity for this purpose. Alternatively, Lee et al. [45] parameterize the likelihood in discrete-ting settings using the probability mass function. Finally, survival clustering methods have been proposed to identify groups of samples based on their survival patterns [7, 33].
also emerged as a flexible approach to predict cumulative hazard functions
tree-ensemble manner [32]. Deep learning-basee
Despite advances in survival analysis, adapting these models for HTE estimation remains understudied. Early investigations focus primarily on binary treatment settings, leveraging tree-based methods such as causal trees [71], Bayesian additive regression trees (BART) [29], and random forests [67] to measure the effect of binary treatments on expected survival time. More recent approaches, such as that by Chapfuwa et al. [12], explore representation learning and generative modeling to estimate event times for different treatments. A seminal work by Curth et al. [14] adopts a multi-task NN
architecture with time- and treatment-speccffee network heads to predict discrete conditional hazard
22

---

functions, incorporating balanced representations to address covariate shifts. Schrod et al. [58] simplify this approach by combining a dichotomous multi-task architecture with Cox- Cox-propottional hazard loss [37]. However, such works focus on binary interventions. In comparison, our DoseSurv model is more versatile and can handle single- and multi-arm treatment regimens with mutually exclusive continuous-valuee (dosage-dependent) and non-dosage treatment options.
D.3 Comparison of DoseSurv with Related Work
Table D.I compares key ML methods for HTE estimation, specifically targeting continuous-valued treatments or TTE outcomes. Unlike related work, DoseSurv uniquely integrates TTE analysis with complex treatment settings that include continuous- mixed treatment scenarios. Due to the lack of HTE estimation models for such settings, our baseline comparisons rely on standard survival models, implemented as model-agnostic T- and S-learners. Although not explicitly studied here, DoseSurv is generally also applicable to treatment scenarios with more than two treatment options.
treatments, non-dosage treatments, and
Table D.I: Overview of selected key neural methods for HTE estimation under continuous-valued treatments or TTE outcomes and their applicability across different settings.
Continuous-valued (Dosage) Treatments
Mixed Dosage and Non-dosage Options
More than Two Treatment Options
TTE Outcomes
Model
DRNet [60]
X
SCIGAN [9]
VCNet [51]
x
SurvITE [14]
X
BITES [58]
X
X
DoseSurv (ours)
E Implementation
E.1
DoseSury
The source code of DoseSurv will be made available at:
Experiments were performed on single GPU machine (specifications: CPU - Intel Xeon W5-2445, EPU-NVIDIA RTX A4000, and RAM 64GB DDR5, OS - Linux). DoseSurv implemented in PyTorc.. The training time of DoseSurv without IPM regularization. Per-epoch training time generally ranges from under a second to a few seconds, depending whether IPM regularization is applied. Overall, computation time increases with the number of IPM bins.
under 30 seconds per for Scenario S1 (n7 0)
Across all main experiments, we implemented DoseSurv with I hidden layer in the representation network and 3 hidden layers each treatment-specffi head of the RBF hazard estimator, each layer comprising 100 nodes. For each layer in the RBF hazard estimator, we use 5 Gaussian RBFs with centers, initialized at {0, 0.25, 0.5, 0.75, 1). The shape parameters of are initialized 0.7. For each experiment, we chose 7 from 10.1, 0.01} based on the lowest INBLL achieved on the validation data. For experiments on synthetic data, we optimized under r 2. Network parameters were optimized using Adam optimizer [38] with learning rate of 0.001. We use ReLU activation functions, implement DoseSurv with batch
dropout probability of 0.1, and a batch size of 500. Moreover, normalization layers, and employ early stopping after 30 epochs without model improvement on the validation data. We adopt the implementation of the finite-sample approximation of the Wasserstein distance from [14]. which in turn follows the implementation in [5. 17], and improve computational efficiency by parallelizing its computation across treatments a, dosage bins q, and time bins t. We use the same parameters, i.e., an entropic regularization strength of Awass = 10 and 10 Sinkhorn iterations. Moreover, choose Q and 5.
Deviating from the primary experiments, we configure DoseSurv for the additional experiment described in C.2 with representation network consisting of 3 layers of 200 nodes each and
23

---

representation size of 50. The treatment- specific hazard estimators are designed with 2 layers of 100 nodes each, without applying batch normalization, Moreover, we choose 2 and 10.
E.2 Baseline Models
We benchmarked DoseSurv against publicly available implementations of four state-of-the ma- chine and deep learning models for survival analysis. DeepSurv DeepHit', RSF and NSC We adopt these implementations as both and S-learners. First, we standardize covariates and treatment-specifa dosages. For the standard (S-learner) baselines, one-hot) encoding of treatment and dosage information. We add the dosage information for each treatment option additional continuous covariates. In scenarios with more than treatment option, we set the dosage values for unobserved treatments to - 10, which does not normally fall within the standardized dosage range, ensuring that the models do not mistakenly interpret this value as a valid dosage. For categorical (non-dosage) treatments, we set the dosage parameter for the observed treatment 1. In addition, for scenarios S2 and S3 with more than one treatment group, we implement T-learner versions of the baseline models: DeepSurv-T. DeepHit-T, RSF-T, and NSC-T. Here, we split the data into the respective treatment groups, and train separate versions of the model on each treatment group individually, including the corresponding dosage as an additional covariate.
slot-based (continuous
We optimize hyperparameters based on grid search. We choose the optimal hyperparameter combination based on the minimum INBLL, which comprises both calibration and discrimina- tive performance. For DeepSurv we optimize the number of layers and nodes, choosing from {1000, [100, 100], [100, 100, 100), [200, 200]]. Additionally, for DeepHit, we optimize the parame- ters Diep {0.2, 0.5) and DeepHit [0.1,1], which control the contribution and properties of the ranking loss. For NSC, we choose network sizes of the mixture weights and survival net- works from [100, 100, 100), [200, 200]]. Furthermore, we choose the number of components for the mixture from {4. 6), and the size of the latent cluster representation from {10. 100). For RSF, we optimize the number of trees {100, 300} as well as the minimum numbers of samples required to split an internal node (f6, 12}) or be at a leaf node ((3, 6)). For all network- based baseline models, as for DoseSurv, we use a dropout rate of 0.1, batch size of 500, batch normalization, and early stopping after 30 epochs.
For experiments on the Twins and breast cancer datasets under binary treatments, we additionally compare against adapted public implementations of the HTE models BITES" and SurvITE ensure compatibility with our environment. For SurvITE, architecture (representation network: [100, 100, 100); hazard estimators: [100, 100).. For BITES, we choose the standard representation network of size [7, 5). and optimize the treatment-speeife heads from {[5, 3). [3]]. As for DoseSurv, we choose the IPM regularization strength for both models from {0.01, 0.1} and BITES's blur parameter from {0.05, 0.13.
choose the standard network
F Remarks on HTE Estimation and Assumptions
In HTE inference, generally aim to estimate
xxx) E Y Yo x =
E Y X = - -
where Yo and Y1 are two continuous potential outcomes under treatments 0 and 1, and X = x are the patient covariates. The potential outcomes can also be expressed using the do-operator [53] indicating
an intervention to made:
=())=EY x
Ey (1= X x
httrp::/ggicaab..omm/maaaa///// "scikit-survival [54]
24

---

For HTE estimation between two dosage-dependent interventions (a. q) and (ā. q) and the time- specific survival probabilities as outcome of interest this translates into
Vsarv (t|3)) P(Ta. X a, do 2 t))
F( > x a, do 2 1))
Sa.g((( (t))
However, since ground-truth survival functions are unknown in real-world settings, we focus on estimating the HTE in terms of the RMST (up until the time horizon, tm) given by
Pmmst (x) ) = [min
(Ta.q: (mminTT.q.
tm) X= -E imin(Ta.q. X=
tan (x)
stust
(x)
2.9
(1)
(a,(()z)
1=
Surr (t/az).
Since (t|2:) = 11 (T|()) therefore need to compute An (T/z) in order to find unbiased estimates for both Usurr (surv(t (8|z)) and Vrmst (z). We can rewrite
\a.g(T/2))
=P(T= 2 2T, x. do a, a 20.0=902)
X xz do T)).
Using Assumptions 1, 3-4, outlined in Section 3.2, we obtain [14]
Aa (T|z))
= 2,, -
= 2x=
= PP0=:T22633==
>(T(2, a, 4)
which can be estimated non- parametrically under the additional assumptions that treatment and dosage assignment (Assumption 2) and the survival and censoring processes (Assumption 5) satisfy positivity (i.e., are non-deterministic)..
G Causal Diagram
The directed acyclic graph (DAG) [53] for our setting is shown in Fig. G.1 and follows prior work on HTE estimation with TTE data [14, 65], extended here to include the dosage Q. We define ((()) 1(T < tt A 1) and (t) 1(T < t, A 0) counting processes indicating, over time, whether an event (death) or, respectively, censoring has occurred. By definition, M and (t)()) are for = 0 and switch to I for / 2 î, i.e., when either the event or censoring occurs at time T'
X consists of (possibly overlapping) subgroups X,, X2, X3, and X₃ of covariates that influence (1) the event process, (2) treatment assignment, (3) dosage assignment, and (4) the censoring process. These subgroups and their causal influences associated with covariate shifts arising from (1) event-induced bias, (2) treatment-selection bias, (3) dosage-selection bias, and (4) censoring bias, respectively [14].
The causal diagram describes how A and Ne evolve over time under the influence of covariates X. treatment A, and dosage Q, and is in line with the Assumptions formulated in Section 3.2. The
25

---

temporal dependencies shown (N 1) + (t) and Ne(t - 1) - Mtt), (c) reflect the fact that these counting processes are sequential an individual's event status at time depends on their status at the previous time step 1 - 1. In other words: once (V 1) = (event occurred by time t t- 1), then (t))= 1 similarly, once (t 1) (censored by time 1), then (t and (tt) 0 (no event can occur after censoring). Furthermore, (t) + (t) captures the logical
constraint that an event precludes censoring- ((t) = then (t)
(( 1)
N(()
X
X
Ne(t
N2(1)
Figure G.I: Assumed DAG for discrete-time survival analysis setting with continuous- treatment, extending prior formulations [14, 65].
26

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We address the problem of HTE estimation for single- and multi-arm treatment regimens with continuous-valuee (dosage-dependent) and non-dosage treatment options, as well as TTE outcomes-a practically important but largely underexplored area. (See Section 3.)
We propose a novel coeffiicoefiicient network with layer-specific RBFs, specifically designed for HTE estimation in the above setting. (See Section 4.)
We introduce computationally efficient IPM regularization approach to address biases in HTE estimation from observational TTE data with continuous-valuee treatments. (See Section 4.)
We validate our model across diverse scenarios and datasets, including reinterpretation of the classical Twins dataset, serving as a real-world benchmark for continuous-value treatments with TTE outcomes. (See Section 5.)
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations of the work in Section 6.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
27

---

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: All assumptions are explicitly stated in Section 3.2. Further mathematical derivations related to estimating the HTE of continuous-valued interventions from TTE data and the corresponding assumptions, are provided in Appendix F.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We present the experimental settings in Section 5 and provide implementation details and hyperparameters for our model and baselines in Appendix E. The data- uuta-generation processes are described in Appendix A. Reviewers can find the source code in the supple- mental material.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
28

---

(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g.. to registered users), but it should possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Reviewers can find the source code the supplemental material. All raw real-world datasets are publicly available. Additionally, the code will be made available GitHub later.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/ public/guides ltaeeeeee/ooeeeeeiiieeeeeeeee www.iiisssoooiiioio..oooo ollicy) for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / /nips. cccccccppooo hte/pullic/pee wtwp//ui///////eee Sttps//ooSSSumiieeeee wwt.uummiissooooooooiiooo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide these details in Appendix E.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
29

---

7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: All main experiments are repeated across multiple runs. We report 95% confidence intervals to capture variability-excepl in Fig. 2 and Fig. 3, where they are omitted for visual clarity.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide the details of computer resources used in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: httpp:::////uuiiieee /nt.riis....//// httpppublcc/chhccccuiieeeeee
Answer: [Yes]
Justification:
We conform to the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
30

---

If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the potential societal impact of our work in Section
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deploymentt of technologics that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA].
Justification: No such risks for our paper.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
31

---

Answer: [Yes]
Justification; We cite the original papers that are used for our code implementation and discuss details in Appendix E.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide implementation details in Appendix E and describe the generation process of the (synthetic) datasets we produce in Appendix A.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA].
Justification:
This paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
32

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]-
Justification: This research does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA].
Justification: The usage of LLMs did not form original, non-standard component of the core method.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non non-standard components.
Please refer to our LLM policy (https: htt////eeippp///// cc/Connferences/2222 for what should or should not be described.
33


Reviews:
[{"review_id": "BVTyc0gntp", "reviewer": "Reviewer_Wk85", "summary": "In this paper, the authors develop a method for estimating treatment effects for survival outcomes with multiple continuous-valued treatments. Following standard assumptions that allow for adjustment for confounding effects, the authors propose a neural network-based architecture that first learns a common, confounding-balanced representation that is fed into treatment-specific event classifiers (classifiers because survival is computed on a time-discretized basis). A novel and interesting aspect of the approach is that continuous dosage is accounted for in the treatment-specific neural network classifiers through RBF-based weights that are a function of dosage. The authors empirically evaluate the model and some ablations on synthetic data, before adapting the common semi-synthetic Twins benchmark to their continuous-value treatment scenario and evaluating vs baselines there. Across evaluations, the proposed approach outperformed baselines.", "questions": "### Identification and Representation balancing\nQ1: An important aspect of the proposed approach is the representation balancing (Section 4.3), because it is how causal estimation is enabled from observational data. In Appendix F it appears the authors provide a standard identification derivation. However, this is not discussed in the context of the IPM regularization approach. Why does IPM regularization preserve identification? When do balanced representations yield unbiased treatment effects?\n\nQ2: The authors use discretized dosage and event times. How does this discretization affect the identification results?\n\nQ3: The IPM regularization scheme is quite sophisticated for the average causal inference practitioner. Did the authors consider using something simpler like generalized propensity scores (see, e.g., Imbens (2000)) which would combine a treatment propensity model with a dose selection model? What about outcome regression with treatment interactions? The authors claim that this is a potential problem in high-dimensional settings (Ln 156). This claim could have been tested and shown in experiments. As it stands, readers don't know if the IPM complexity is necessary or if simpler approaches would work just as well.\n- Imbens, G. W. (2000). The role of the propensity score in estimating dose-response functions. Biometrika, 87(3), 706-710.\n\n### RBF Hazard Estimator\nQ4: To my understanding, the RBF Hazard estimator involves a potentially very large number of parameters. Further, the authors loosely refer to \"sufficient data\" (Ln 170) enabling the flexibility of this approach. Given that these sub-models are fit only on samples which received a particular treatment, can the authors speak to the sample size requirements or convergence properties of their approach? This is helpful for understanding which applications might have sufficient data for using the proposed approach.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Quality/Significance:\n- Overall, I think the quality of this manuscript is good. The authors address important aspects of treatment effect estimation (survival outcomes & continuous treatments) that are often ignored but are practically relevant. The approach is interesting, and I think the RBF-hazard estimator is likely to be of interest to other researchers in this area.\n- Evaluation is hard for causal inference because in true real-world scenarios we do not have ground truth counterfactual outcomes. Thus, the synthetic and semi-synthetic experiments are reasonable. I think it would have been interesting to see the authors illustrate their approach on an actual real-world scenario with multiple, continuous-valued treatments to see how the potential insights that could be gained might differ from prior methods or methods that cannot account (as easily) for different dosages. However, this is not a significant weakness of the paper, though I think it would have made the potential for practical impact much higher.\n- In the questions section I get at a few areas of potential weakness, mainly identification assumptions, sample size requirements, and some alternative design choices\n\n### Clarity:\n- The manuscript is well-written. The problem is motivated and well-explained and the approach in Section 4 walks through the model in an understandable way.\n- The experiments section in the main paper focuses on the set-up for the experiments, with limited discussion of the actual results. Part of this is because the evaluation is quantitative and metric-based, so there is not much to say when performance is better. That said, it would have been nice to have the results discussed in a way that highlights potential impact of the authors' method. For example, it is my understanding that dose-response curves can be computed as a by-product of the proposed approach. This could be very interesting to visualize and walk through in a qualitative discussion on an actual real-world example.\n\n### Originality:\n- I am not up to date on the latest methods for causal effect estimation with survival outcomes and continuous-valued treatments. That said, the authors approach is interesting and reasonable, and I think others working on these problems will find aspects of the proposed method useful in their own work.", "was_revised": true, "final_justification": "The authors provided an in depth response to all of my questions and concerns. Thus, I raised my score to 5 (accept).", "rebuttal": "Dear Reviewer, we appreciate your constructive feedback and are happy to hear that you liked the quality of our paper. Below, we address your questions and comments:\n\n\n## **Replies to Questions:**\n### **Q1: IPM Regularization and Identification**\n\n- Theoretical identification is established through Assumptions 1-6, with derivations provided in Appendix F. These assumptions define the conditions under which unbiased treatment effect estimation is possible. IPM regularization does not \"preserve\" or break identification; rather, it helps achieve good finite-sample performance under these assumptions. \n\n- This is because, even when identification assumptions hold, finite samples can lead to biased estimates in practice due to covariate imbalance between treatment groups. Our IPM regularization scheme addresses this issue by encouraging domain-invariant representations across treatments, dosages, and time points.\n\n- Balanced representations yield less biased treatment effects because they reduce the model's reliance on spurious correlations that arise from systematic differences between treatment populations. This is particularly important in survival data, where time-varying covariate shifts can compound biases.\n\n- Our approach extends IPM balancing from binary treatments [1,2] to continuous treatments with survival outcomes, where we must handle multiple sources of bias, including treatment and dosage selection bias, as well as event- and censoring-induced shifts. \n\n### **Q2: Effect of Discretization**\n\n- Our method involves two distinct discretizations: \n  - (1) We discretize time into intervals for hazard estimation, which is standard practice in discrete-time survival models.   \n  - (2) For computing the IPM regularization loss, we discretize dosage ($q \\rightarrow \\mathfrak{q}$) and time ($t \\rightarrow \\mathfrak{t}$, where $\\mathfrak{t}$ is more coarsely discretized than $t$) solely to enable efficient Wasserstein distance computation between two empirical distributions.\n\n- Our ablation studies (Appendix C.3) demonstrate that DoseSurv remains relatively robust across different time discretizations in the hazard estimator. We have chosen 30 time discretizations by default, as this is the most common discretization in the literature (see DeepHit or SurvITE).\n\n- Importantly, the RBF hazard estimator maintains continuous dose-response modeling—discretization affects only the balancing computation, not the core treatment effect estimation.\n\n- Neither discretization affects our theoretical identification assumptions (1-6), which remain valid for the continuous case. The IPM discretization is purely computational, analogous to histogram approximations in optimal transport, and does not impair our general assumptions, as discussed in our reply to Q1. \n\n### **Q3: Alternative Approaches to Address Confounding**\n- **GPS vs IPM-based Approaches**: Traditional GPS-based reweighting addresses confounding in the initial treatment and dosage assignment at time $t = 0$, but it cannot account for time-varying biases specific to survival data—such as censoring- and event-induced shifts. In contrast, our IPM-regularization approach balances representations across time points, treatments, and dosages simultaneously. \n\n    In theory, one could adapt the idea of reweighting observations for a dynamic survival setting (as outlined by [2] for binary treatments) and combine it with the IPM-regularization approach. In our setting with continuous-valued treatments, the optimal importance weights for reweighting the empirical risk by the marginal covariate distribution $P_0(x) = \\mathbb{P}(X = x)$ \n    and the observational covariate distribution \n    $P_{a,q,\\tau}(x) = \\mathbb{P}(X = x \\mid A = a, Q = q, \\tilde{T} \\ge \\tau)$ \n    are given by: \n\n    $$w_{a,q,\\tau}(x)=\\frac{P_0(x)}{P_{a,q,\\tau}(x)}=\\frac{\\mathbb{P}(A=a, Q=q, \\tilde{T} \\ge \\tau)}{\\mathbb{P}(A=a, Q=q \\mid  X=x) \\cdot \\mathbb{P}(\\tilde{T} \\ge \\tau \\mid A=a, Q=q, X=x)}$$\n\n    where: \n    - $ \\mathbb{P}(A=a, Q=q \\mid X=x)$ is the generalized propensity score for treatment $a$ and continuous dosage $q$ given covariates $x$,\n    - $\\mathbb{P}(\\tilde{T} \\ge \\tau \\mid A=a, Q=q, X=x)$ is the conditional survival probability at time $\\tau$, \n    - $\\mathbb{P}(A=a, Q=q, \\tilde{T} \\ge \\tau)$ is the joint probability of treatment $a$, dosage $q$, with the event $\\{\\tilde{T} \\ge \\tau\\}$.\n\n\n    However, these weights are unknown and must be estimated. In a survival setting, this leads to a significantly more complex framework than the more straightforward IPM-based approach. Furthermore, estimating these weights often results in extreme values—especially at later time points where fewer samples remain—making training unstable. In principle, one could also approximate the weights $w_{a,q,\\tau}$ imperfectly using the empirical frequencies from the time-, treatment-, and (discretized) dosage-specific at-risk populations. However, our additional experiments with reweighted loss terms did not yield improved performance (see our replies under _\"At-Risk Imbalance in Loss Formulation\"_ to Reviewer Bytv and under _\"Censoring in Loss Function\"_ to Reviewer Lefa). For simplicity, we therefore refrain from applying additional weights in DoseSurv’s loss function and instead rely on the IPM-based regularization approach, which handles multiple sources of bias in a unified and stable manner.\n\n- **Outcome regression with interactions and high-dimensional settings**: Our baselines represent survival-adapted ML approaches (neural, tree-based, clustering-based) that treat dosage as a regular feature with potential interaction terms. However, we do show in our experiments that when dosage is treated as just another covariate among many (20-45 features in our experiments), these models struggle to capture the specialized dose-response relationships. This leads to degraded dose-response estimation and consistent underperformance relative to DoseSurv (see results in Fig. 2 and Table 1), which explicitly models the dose-response function through RBF parameterization, across both synthetic and real-world benchmarks.\n\n### **Q4: Sample Size Requirements and Convergence**\n\n- The \"sufficient data\" reference relates to the theoretical universal approximation properties of RBFs in reproducing kernel Hilbert spaces, not specific sample size requirements for our method.\n\n- Our sample size ablation (Appendix C.1) shows DoseSurv achieves competitive performance across all sample sizes (2K-18K), with best performance except at the smallest size (2K), indicating reasonable but not minimal sample requirements. However, it is important to highlight that, aside from non-medical uses that typically involve larger sample sizes, the primary clinical goal of our model is to leverage the extensive, albeit potentially biased, information found in large observational datasets, as opposed to relying on smaller-scale randomized clinical trials, where sample size can indeed become a limiting factor.\n\n- RBF parameterization is actually more parameter efficient than alternative basis functions—our K=5 RBFs compare favorably to B-splines, which require K=8 functions (Ablation A2) while maintaining expressiveness (see Sections 5.1.2 and 5.5.2). \n\n## **Replies to Other Comments:**\n\n- We agree that incorporating more real-world examples would further enhance the demonstration of practical impact. However, it is important to note that publicly available datasets containing continuous dosage information are extremely limited, making such validation challenging in practice. Moreover, evaluating causal inference models on observed (factual) outcomes is generally insufficient, as these outcomes are subject to the very biases our approach is designed to address. This limitation makes it difficult to draw meaningful conclusions from factual-only evaluations. In comparison to most existing work on HTE with continuous treatments, which typically relies exclusively on synthetic data (e.g., [3,4,5]), our study already offers a notably broader real-world assessment, which includes an adaptation of the Twins dataset. Additionally, we present results for a binary treatment case involving hormone therapy (Appendix C.2) to illustrate how DoseSurv could be applied to analogous clinical settings involving dosage-dependent treatments. \n\n- Regarding dose-response visualization: We use RMST as the outcome of interest and compute predicted dose-response curves by evaluating the trained DoseSurv model under different dosages. We visualize ground-truth dose-response curves in Appendix Figure A.1 and are happy to provide further visualizations of predicted dose-response curves and corresponding discussions in the final version of our paper.\n\n**References:**\n\n[1] U. Shalit et al., \"Estimating individual treatment effect: Generalization bounds and algorithms,\" in Proc. 34th ICML, vol. 70, pp. 3076–3085, Aug. 2017.\n\n[2] A. Curth et al., \"SurvITE: Learning heterogeneous treatment effects from time-to-event data,\" in Adv. Neural Inf. Process. Syst., vol. 34, pp. 26740–26753, 2021.\n\n[3] P. Schwab et al., \"Learning counterfactual representations for estimating individual dose-response curves,\" in Proc. AAAI Conf. Artif. Intell., vol. 34, no. 4, pp. 4046–4053, Feb. 2020.\n\n[4] I. Bica et al., \"Estimating the effects of continuous-valued interventions using generative adversarial networks,\" in Adv. Neural Inf. Process. Syst., vol. 33, pp. 16434–16445, 2020.\n\n[5] L. Nie et al., \"Varying coefficient neural network with functional targeted regularization for estimating continuous treatment effects,\" in Proc. ICLR., 2021."}, {"review_id": "uF0EkSN1lK", "reviewer": "Reviewer_Lefa", "summary": "The authors introduce DoseSurv, a method for continuous-valued treatment effect estimation in the survival setting. DoseSurv employs a shared-representation scheme based on the method of integral probability metric regularization for effect estimation, and a radial basis function hazard estimator to estimate survival and dose response curves. The method outperforms T-learner baselines of other survival methods in the large-sample regime on synthetic and empirical data.", "questions": "1. Could the authors separately benchmark their RBF hazard estimator (described in the third bullet point under Weaknesses, above)? [This would contribute to me raising my Quality score to a \"3: good\"]\n2. Could the authors provide additional results on a dose-response dataset, preferably with moderate-to-few samples, and benchmark against an RBF-based T-learner baseline? [This would also contribute to me raising my Quality score to a \"3: good\"; if experiments are extensive, I would raise my Quality score to a \"4: excellent\"]\n3. Could the authors more cleanly characterize the novel elements of this work, versus those that use adopted technical machinery? [Depending on the response, I can likely raise my Originality score to a \"3: good\"]", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n* The paper is very clear, thoughtful, and well-written. It was a pleasure to read -- I thank the authors for the opportunity to review their work.\n* I enjoyed the problem setup -- the notion of continuous-valued treatments for dosage response is sensible and clearly explained.\n* The baselines are strong. I especially appreciated how the authors benchmarked against T-learner variants of each model, rather than the observational version of the model.\n\nWeaknesses:\n* The Empirical Risk Loss equation (Line 180) does not appear to account for censoring. If so, this could entail significant bias in the estimator. How is censoring accommodated in the loss function when learning the model?\n* Novelty: the work could do more to clearly characterize the elements of the work that are novel, versus that which relies on technical machinery developed from other sources. (a) in Section 4, I would have liked to see a citation to [1]. (b) I would have also liked to see references to the literature on discretizing continuous-valued treatments as a means of simplifying the effect estimation problem. (c) I would have further liked to see references to the literature on radial basis function (RBF) function estimation, especially if prior work has used RBFs for survival estimation.\n* The methodology of the paper makes two key contributions: (a) subgrouping and IPM to obtain continuous-valued dose-response estimates, and (b) the use of an RBF estimator of the hazard. I wish the authors had done more to benchmark the two contributions sequentially -- first (b), then (a).\n  1. To benchmark (b) -- ignoring the task of estimating the individual dose-response curve, does an RBF hazard estimator compare favourably to the standard neural survival analysis baselines? (DeepSurv, DeepHIT, etc.) If not, is the RBF hazard estimator component of DoseSurv exchangeable with one of these other architectures? Why or why not?\n  2. Benchmarking (b) in the observational setting justifies your design decision to use an RBF hazard estimator in DoseSurv. Then, your existing experiments serve to benchmark the utility of subgrouping and IPM to obtain continuous-valued dose-response estimates. (Although for this component, I would appreciate if you would add a benchmark to a T-learner variant of the RBF hazard estimator).\n* The authors do little to justify their design decisions -- why RBFs as opposed to any other estimator? Why discretize as the means of learning the continuous-valued treatment effect? Even without rigorous experiments to justify their design decisions in the space of all possible decisions, I would have appreciated more justification of these choices (so that, say, a future effort building on this work knows what choices were made in its development).\n* The method appears to only work in very large-data regimes -- 20,000 instances in the synthetic data, and >5,000 instances in the empirical data. How does the method perform in the more realistic smaller-sample settings that are common in healthcare dosing data? (E.g., in a single-site or clinical trial dataset).\n* There is mild misalignment between the presented experiments and the empirical use case that the authors detail. I would appreciate if they authors could provide additional results on a dose-response dataset (rather than using twin weight at birth as a synthetic dose), such as MIMIC-IV.\n* Minor nit: in Equation 2, should the $C \\geq \\tau$ be outside of the $\\textbf{do}$ operation?\n\n[1] Shalit, Uri, Fredrik D. Johansson, and David Sontag. \"Estimating individual treatment effect: generalization bounds and algorithms.\" In International conference on machine learning, pp. 3076-3085. PMLR, 2017.", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer, thank you for your thorough review and feedback. Below, we address your main questions and comments:\n\n## **Replies to Questions:**\n\n### **Q1: Additional Benchmark Experiments**\n\nWe understand you would like us to benchmark separately (a) our IPM regularization scheme and (b) the RBF hazard estimator, which uses varying-coefficient networks with radial basis functions. \n\n- As requested, we provide performance results for our RBF-based varying coefficient hazard estimator, DoseSurv's core component, without the representation network or IPM components. The following results are from the single-treatment scenario S1 ($\\eta=2$) across 5 runs:\n\n    ||$\\sqrt{\\text{MISE}}$|$\\sqrt{\\text{DPE}}$|$C^{\\text{td}}$|\n    |-|-|-|-|\n    |DeepSurv|7.821 ± 0.046|10.618 ± 0.049|0.639 ± 0.003|\n    |DeepHit|7.859 ± 0.017|10.437 ± 0.052|0.629 ± 0.006|\n    |RSF|7.672 ± 0.010|9.367 ± 0.036|0.574 ± 0.005|\n    |NSC|7.462 ± 1.882|9.497 ± 1.522|0.623 ± 0.106|\n    |DoseSurv (no IPM)|5.723 ± 0.264|8.206 ± 0.242|0.723 ± 0.005|\n    |DoseSurv|4.708 ± 0.383|8.049 ± 0.237|0.744 ± 0.005|\n    |RBF|6.567 ± 0.205|8.599 ± 0.155|0.699 ± 0.003|\n\n   The RBF hazard estimator alone achieves performance similar to DoseSurv without IPM regularization. This is expected in this single-treatment scenario, as DoseSurv reduces to a single RBF hazard estimator with extra linear layers from the representation network. While the RBF-based hazard estimator is outperformed by the full DoseSurv framework with IPM regularization, it still outperforms all baseline models. This demonstrates the superiority of the varying-coefficient design with RBFs for continuous-valued treatments.\n\n\n- You asked if the RBF hazard estimator in DoseSurv can be replaced with other architectures. Below, we provide results for DoseSurv under Scenario S2 ($\\eta=2$), where the RBF hazard estimator was replaced with an MLP (hazard estimator: 2 hidden layers; representation network: 4 hidden layers; 100 nodes each). We also provide results for our standard DoseSurv for comparison:\n\n    ||$\\sqrt{\\text{MISE}}$|$\\sqrt{\\text{DPE}}$|$C^{\\text{td}}$|\n    |-|-|-|-|\n    |DoseSurv MLP without IPM| 7.249 ± 0.209|9.308 ± 0.232 |0.593 ± 0.007|\n    |DoseSurv MLP with IPM|6.920 ± 0.531|7.993 ± 1.012|0.626 ± 0.017|\n    |Standard DoseSurv (RBF with IPM)|5.164 ± 0.185|5.852 ± 0.319| 0.668 ± 0.007|\n\n    The results indicate our IPM regularization scheme may also be beneficial with different hazard estimator designs, beyond DoseSurv's standard RBF-based varying coefficient network.\n\n### **Q2: RBF-based T-Learner**\n- As requested, we provide performance results for an RBF-based varying coefficient hazard estimator, implemented as a T-learner (RBF-T), under Scenario S2 ($\\eta=2$) across 5 runs. For comparison, we report the performance of DoseSurv and all T-learner baselines. We find that the performance of RBF-T lies between that of DoseSurv and the baselines:\n\n    ||$\\sqrt{\\text{MISE}}$|$\\sqrt{\\text{DPE}}$|$C^{\\text{td}}$|\n    |-|-|-|-|\n    |DeepSurv-T|6.998 ± 0.040|9.201 ± 0.819|0.604 ± 0.002|\n    |DeepHit-T|6.940 ± 0.009|9.402 ± 0.177|0.605 ± 0.002|\n    |RSF-T|6.736 ± 0.026|7.547 ± 0.076|0.610 ± 0.008|\n    |NSC-T|7.125 ± 0.495|9.397 ± 0.463|0.599 ± 0.019|\n    |DoseSurv|5.164 ± 0.185|5.852 ± 0.319| 0.668 ± 0.007|\n    |RBF-T|6.113 ± 0.324|7.070 ± 0.362|0.645 ± 0.009|\n\n- We acknowledge that the varying-coefficient approach in the hazard estimator requires sufficient data to accurately model the complex dose-response relationships in our datasets. This is why we provide an ablation study on sample size in Section C.1 of the Appendix. Below, we present additional results for RBF-T with reduced complexity (1 hidden layer, 50 nodes) in a low sample-size regime (300 training, 50 validation samples) under Scenario S2 ($\\eta=2$). We compare this with other non-parametric neural T-learners—NSC-T and DeepHit-T. We find that all models achieve similarly low performance across 5 runs:\n\n    |Model|$\\sqrt{\\text{MISE}}$|$\\sqrt{\\text{DPE}}$|$C^{\\text{td}}$|\n    |-|-|-|-|\n    |DeepHit-T|7.550 ± 0.232|9.670 ± 0.097|0.540 ± 0.017|\n    |NSC-T|8.831 ± 0.589|9.596 ± 0.106|0.541 ± 0.013|\n    |RBF-T|8.293 ± 0.081|9.694 ± 0.018|0.555 ± 0.006|\n\n    However, apart from non-medical applications, which often involve much larger sample sizes, the primary clinical motivation behind our model is to leverage the rich, yet potentially biased, information in large observational datasets, rather than relying on small-scale (randomized) clinical trials.\n\n### **Q3: Clarification of Novelty**\n\n- Our main contribution lies in addressing the highly relevant yet understudied problem of estimating HTEs of dosage-dependent treatments from survival data—a setting where specialized baselines, performance metrics, and benchmark datasets are either lacking or require adaptation. Previous ML research on continuous treatments has mainly focused on continuous or binary outcomes only.\n\n- We introduce a systematic framework that integrates carefully designed components to establish a foundation for this novel setting. Specifically, we propose a new model, DoseSurv, which employs varying coefficient networks with layer-specific Gaussian RBFs in its hazard estimator. While neural varying-coefficient models have been introduced previously for estimating HTEs for continuous outcomes [1], we adapt this concept for survival settings and introduce a novel approach to modeling the network parameters using Gaussian RBFs, which promise universal approximation properties and long-range support across the full dosage range, motivated by the need to handle regions with limited overlap.\n\n- IPM regularization using Wasserstein distance was introduced for binary treatment settings [2] and has also been adopted for survival data [3]. Our contribution is to extend these frameworks to continuous-valued treatment settings combined with survival data; the technical machinery to implement our approach was adapted from [3].\n\n- Finally, we propose using the well-established Twins dataset as a real-world benchmark for continuous-valued treatment settings with TTE data, treating birthweight as a continuous-valued dosage. Additionally, we create synthetic TTE datasets for our scenarios.\n\n    We are convinced that both the causal ML and survival analysis communities would greatly benefit from our in-depth exploration of this widely overlooked yet highly important problem.\n\n\n## **Replies to Other Comments**\n\n### **Censoring in Loss Function**\n- The loss term in Eq. 7 uses the summation $\\sum_{i:\\tilde{\\tau}_i \\ge t}$, ensuring only at-risk patients contribute at each time $t$, thereby correctly handling censoring in discrete-time survival.\n\n- As explained to Reviewer Bytv, it would generally be possible to reweight each sample's loss contribution by the number of samples in the respective time-, treatment- and (discretized) dosage-specific at-risk population. However, in practice, this is not always beneficial, as, for example, noise from sparse later time points or less frequent treatment regimens may be amplified. We found that applying such weights often actually worsened performance. For example, in S1 ($\\eta=0$), our standard DoseSurv model achieved slightly better performance across 5 runs, compared with the same model with inversely weighted loss terms. Therefore, we prioritized simplicity and refrained from reweighting individual loss terms.\n\n    ||$\\sqrt{\\text{MISE}}$|$\\sqrt{\\text{DPE}}$| \n    |-|-|-|\n    |DoseSurv (unweighted)|4.21 ± 0.25|7.88 ± 0.24|\n    |DoseSurv (with sample weights)|4.45 ± 0.27|8.09 ± 0.17| \n\n### **Design Justifications**\n- As explained above, we use Gaussian RBFs due to their favorable properties, including universal approximation, long-range support, learnable centers and shapes, and parameter efficiency. We provide an ablation study (see Sections 5.1.2 and 5.2.2) to justify our design. Discretization of dosages and times is necessary during IPM regularization for estimating the Wasserstein distance between two empirical distributions; however, dosages are treated as continuous for the actual HTE estimation. We are happy to extend these design justifications in Sections 4.2 and 4.3 for the final version of our paper.\n\n### **Real-world Data**\n- Evaluating on additional real-world datasets is not feasible, as counterfactual outcomes are generally unavailable outside the Twins dataset. Assessing performance on factual outcomes alone is typically insufficient due to the inherent biases our model is designed to address. Moreover, publicly available dosage-survival datasets are limited. Compared to prior work on HTE with continuous treatments—which often relies solely on synthetic data (e.g., DRNet, SCIGAN, VCNet)—our evaluation, in fact, represents a more thorough use of real-world data. Additionally, we present results from a real-world clinical setting with binary treatments in Appendix C.2 to illustrate how DoseSurv could be employed in a corresponding dosage-dependent application. \n\n### **Explanation of $do(..., C\\geq\\tau)$ in Eq. 2**\n- This is not a mistake: To estimate a causal effect, we must implicitly treat censoring as an action we can intervene on, which means we require samples to be 'set' to uncensored [3,4]—a formal intervention justified by the Independent Censoring Assumption.\n\n\n**References:**\n\n[1] L. Nie et al., \"Varying coefficient neural network with functional targeted regularization for estimating continuous treatment effects,\" in Proc. ICLR., 2021.\n\n[2] U. Shalit et al., \"Estimating individual treatment effect: Generalization bounds and algorithms,\" in Proc. 34th ICML, vol. 70, pp. 3076–3085, Aug. 2017.\n\n[3] A. Curth et al., \"SurvITE: Learning heterogeneous treatment effects from time-to-event data,\" in Adv. Neural Inf. Process. Syst., vol. 34, pp. 26740–26753, 2021.\n\n[4] O. M. Stitelman and M. J. van der Laan, \"Collaborative targeted maximum likelihood for time to event data\", Int. J. Biostat., vol. 6, no. 1, p. 12, Jan. 2010."}, {"review_id": "s8ikaTgVij", "reviewer": "Reviewer_Bytv", "summary": "This paper presents DoseSurv, a versatile method for estimating HTE in survival analysis. DoseSurv generalizes to multiple treatment types, including those with continuous dosages. It achieves this by employing a deep learning network that models the hazard function based on time, covariates, and specific treatment pairs (type and dosage). To effectively incorporate continuous dosages, the network parameters are designed as a continuous function of dosage levels, $q$, utilizing learnable RBF kernels. This design helps preserve the influence of dosage levels, even with relatively high-dimensional covariate inputs. Furthermore, the method mitigates selection bias by balancing representations through an IPM term, which is applied by discretizing continuous dosages into equally sized subgroups.", "questions": "Regarding Weakness 1: The varying number of \"at risk\" patients across different time points can introduce exposure biases, particularly toward earlier time points. This is because, in time-to-event data, the number of events significantly decreases at later time points. As a result, the \"long format\" loss function (7) will reflect this imbalance, having considerably fewer observations for those later times. To address this, a better approach might involve either using an MLE approach (as in [A], [B]) or incorporating balancing representations if the \"long format\" is used (as in SurvITE). Please elaborate on the underlying design choice of the proposed method in light of these challenges.\nRegarding Weakness 2: The real-world experiments are limited to the TWIN dataset. Expanding the empirical validation to include other real-world datasets and assessing the results against domain expertise would considerably strengthen the practical implications of this work.\nRegarding Weakness 3: How are the HTE-sensitive parameters selected using the validation set when the evaluation metrics are not capable of assessing the HTE given observational data?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\nThe paper is well organized and clearly written.\nThe paper introduces a novel and well-justified approach by making the network parameters a function of dosage levels.\nWeakness:\nThe paper doesn't clearly discuss two main design choices for the treatment-specific hazard estimators: (i) The authors use a multi-task network for binary predictions across the entire time horizon. However, an alternative approach—feeding time ($t$) as an input to a shared network (similar to SurvITE, DRSA [A], and ConSurv [B])—could help manage a large number of outputs. The reasoning behind their chosen design isn't explained. (ii) The paper also doesn't discuss the implications of training the hazard estimator using MLE, as seen in [A] and [B], without using a \"long format\" for the data. Employing the long format could introduce additional bias, as the number of patients \"at risk\" can vary significantly between earlier and later time points (as pointed out by the authors in SurvITE).\nThe real-world experiments are limited to a single semi-synthetic dataset.\nThe authors do not clearly explain how parameters sensitive to HTE—like the number of discretization bins or loss coefficients—were chosen using the validation set. This is particularly concerning because INBLL doesn't directly measure the accuracy of treatment effect estimation.\n[A] Ken et al., “Deep Recurrent Survival Analysis,” AAAI 2019.\n[B] Lee et al., “Toward a Well-Calibrated Discrimination via Survival Outcome-Aware Contrastive Learning,” NeurIPS 2024.", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer, thank you for your thoughtful review. We appreciate your technical insights and address your specific concerns below:\n\n### **Reasoning Behind Design Choices in our Survival Model**\n\n- Our treatment-specific network heads adhere to the same standard network design found in other classification-based approaches, such as DeepHit: a single network shared across all discrete time points, with separate outputs for each time point.\n\n- In fact, SurvITE does not utilize a shared network with a single output across time points. Instead, it employs individual time-specific network heads, which significantly increases complexity compared to our approach. Notably, SurvITE's appendix analysis demonstrates that simply feeding time as input to a shared network head (as you suggested) may actually worsen performance—particularly when ample training data is available—compared to using separate outputs for each time point. Hence, we did not follow this approach.\n\n- While MLE-based survival models like DRSA and ConSurv are effective for likelihood-based TTE modeling, we opted for a classification-based approach due to its simplicity and compatibility with flexible, time-varying treatment effect estimation—especially when employing varying coefficient architectures for continuous treatments. In these cases, chaining hazards through a survival likelihood can complicate optimization and reduce stability.\n\n### **At-Risk Imbalance in Loss Formulation**\n\n- You correctly identify that the varying number of at-risk patients could theoretically be addressed by inverse-weighting each loss term according to the number of samples in each time-specific at-risk group. However, one drawback of this approach is that it may amplify noise from sparse later time points or less frequent treatment regimens, potentially creating training instability due to varying loss magnitudes across time points, treatments, and dosages.\n\n- In fact, we conducted experiments where each sample's loss contribution was inversely reweighted by the number of samples in the respective time-, treatment-, and (discretized) dosage-specific at-risk population. We found that applying such weights led to minor changes in performance, often negative. For example, in S1 ($\\eta=0$), our standard DoseSurv model achieved slightly better performance across five runs compared to the same model with inversely weighted loss terms. Therefore, we prioritized simplicity and refrained from reweighting individual loss terms.\n\n    |  | $\\sqrt{\\text{MISE}}$ | $\\sqrt{\\text{DPE}}$ | \n    |--|--|--|\n    | DoseSurv (unweighted) | 4.21 ± 0.25 | 7.88 ± 0.24|\n    | DoseSurv (with sample weights) | 4.45 ± 0.27 | 8.09 ± 0.17| \n\n  Please also see our reply to Reviewer Wk85 under question Q3, where we expand on reweighting approaches in HTE survival models and provide an expression for the optimal importance weights (involving generalized propensity scores) needed for optimal reweighting in survival settings under continuous-valued treatments.\n\n\n### **Real-World Validation**\n\n- We acknowledge that further validation on real-world datasets would always be desirable, but it is practically difficult due to the absence of counterfactual outcomes in datasets other than Twins and the general scarcity of publicly available datasets with continuous dosage information. Evaluating causal models on factual outcomes is generally inadequate because of the underlying biases in the test data—precisely the biases our model aims to overcome. Demonstrating this capability on factual outcomes would not be feasible. Furthermore, compared to other HTE literature with continuous-valued treatments, which relies purely on synthetic data (e.g., DRNet [1], SCIGAN [2], VCNet [3]), our evaluation is, in fact, already more extensive with respect to real-world datasets. Additionally, we include results from a real-world clinical setting with binary treatments in Appendix C.2; specifically, DoseSurv was applied to hormone therapy decisions for breast cancer patients and showed strong predictive performance. A similar approach could be extended to settings with dosage-dependent treatments.\n\n\n### **HTE Parameter Selection Strategy**\n\n- As we describe in the Limitations paragraph (l. 333-337), selecting hyperparameters in causal settings from validation data comprising only factual outcomes can indeed be challenging and may lead to biases. For this reason, we keep most of DoseSurv's model and training hyperparameters fixed across the main experiments, demonstrating that our choice is robust across different settings.\n\n- In particular, we do not change or optimize the number of discretization steps across our main experiments, setting it to 30, which is the most common discretization in the literature (see DeepHit or SurvITE). However, we also report results under different time discretizations in Appendix C.3.\n\n- As described in Appendix E.1, we optimize the regularization strength from {0.01, 0.1} based on the lowest INBLL on the validation set. Alternatively, we describe a more advanced approach, which would involve selecting the highest value of $\\gamma$ that does not impair INBLL on the validation set, balancing regularization and predictive power.\n\n- You are correct in pointing out that INBLL, like any other metric we can compute in real-world scenarios, does not directly measure the accuracy of HTE estimates. This is precisely the fundamental problem of causal inference, which requires us to select hyperparameters based solely on factual outcomes. INBLL is a suitable metric for survival settings as it captures both discrimination and calibration performance [4].\n\n\n**References:**\n\n[1] P. Schwab et al., \"Learning counterfactual representations for estimating individual dose-response curves,\" in Proc. AAAI Conf. Artif. Intell., vol. 34, no. 4, pp. 4046–4053, Feb. 2020.\n\n[2] I. Bica et al., \"Estimating the effects of continuous-valued interventions using generative adversarial networks,\" in Adv. Neural Inf. Process. Syst., vol. 33, pp. 16434–16445, 2020.\n\n[3] L. Nie et al., \"Varying coefficient neural network with functional targeted regularization for estimating continuous treatment effects,\" in Proc. ICLR., 2021.\n\n[4] H. Kvamme, Ø. Borgan, and I. Scheel, \"Time-to-event prediction with neural networks and Cox regression,\" J. Mach. Learn. Res., vol. 20, no. 129, pp. 1–30, 2019."}, {"review_id": "swsPgXuwNq", "reviewer": "Reviewer_3uNJ", "summary": "This paper proposes DoseSurv, a method for estimating personalized dose-response curves from censored observational survival data. By integrating censoring-aware loss and adversarial balancing, it addresses confounding and censoring simultaneously. Experiments show improved survival prediction and dose recommendation over existing baselines.", "questions": "1. What is the meaning of the last paragraph in Section 3.1? It contains a lot of information, but it’s unclear how it connects to the following sections.\n\n2. In the annotation of Figure 1, why is custom discretization of dosage necessary? Could this cause information loss?\n\n3. How do you ensure the overlap assumption holds after representation learning?\n\n4. Could you provide more details on how the RBF hazard estimator works? What about using other kernel functions?\n\n5. In Section 4.3 (Empirical Risk Loss), how should $h(t)$ be chosen in practice?\n\n6. For the experiments, have you considered varying the degree of dosage bias and fixed treatment selection bias?\n\n7. In Table 1, how are the confidence intervals computed? If results are based on only 5 runs, are the CIs reliable?\n\n**Typo**:  \n- Line 96: \"whther either\" is duplicated.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**:\n- **Clarity and Accessibility**: The paper is very clearly written and easy to follow. \n\n- **Strong Motivation**: The motivation is compelling, and the problem of handling (multiple) continuous-valued treatments with time-to-event (TTE) outcomes is both important and timely. The paper makes a strong case for developing new methods in this area.\n\n**Weakness**:\n- **Unfamiliar Treatment Setup**: The setup involving treatment arms and dosage levels is quite new to me. It would be helpful if the authors could provide more references or background on this setting.\n\n- **Missing Causal Diagram**: A causal graph illustrating the assumed data-generating process would greatly clarify the structure and assumptions behind the method.\n\n- **Lengthy Data Description**: The description of the real-world dataset is overly detailed in the main text. Moving some of it to the appendix could improve the flow of the experiment section.", "was_revised": true, "final_justification": "The authors clearly explain all the concerns that I have, and especially the part of the causal graph is very clear.", "rebuttal": "Dear Reviewer, thank you for your valuable feedback. We address your questions and concerns below:\n\n## **Replies to Questions:**\n\n### **Q1: Meaning of Last Paragraph in Section 3.1**\n\n- The final paragraph of Section 3.1 establishes the person-period transformation that enables our discrete-time hazard approach, directly connecting to Section 4's neural network formulation. It explains why we can treat survival estimation as sequential binary classification problems—this is fundamental to understanding our model architecture.\n\n\n### **Q2: Custom Discretization Rationale**\n- The discretization of the dosage, $q$, is solely for IPM regularization and is necessary for estimating the Wasserstein distance between two empirical distributions. The RBF estimator continues to model dosage as a continuous variable.\n\n- No information loss occurs because: (1) discretization is only for balancing purposes, (2) RBF parameters remain continuous functions of $q$, and (3) the continuous $q$ is included as an additional feature in the hazard estimator.\n\n### **Q3: Overlap Assumption and Representation Learning**\n- IPM regularization in representation learning preserves overlap by aligning the treated and control distributions in feature space, trading off predictive power for improved balance (i.e., lower bias) in the learned representation.\n\n- IPM enhances conditional overlap in the learned embedding space for better generalization; however, it does not create or resolve true violations of overlap in the data.\n\n\n### **Q4: RBF & Alternative Basis Functions**\n- The RBF hazard estimator models each network parameter as a continuous function of dosage $q$ using linear combinations of basis functions, which enables smooth dose-response relationships, as detailed in Section 4.2. We propose using Gaussian RBFs for this purpose due to their universal approximation capabilities in RKHS, their long-range support properties, and relatively low complexity.\n\n- In addition to RBF kernels, we also test the varying coefficient network approach with polynomial bases, as described in Section 5.1.2.\n\n### **Q5: Choosing $h$ in the Empirical Risk Loss**\n\n- $h_t^{(a_i)}(\\Phi(x_i,q_i))$ in the empirical risk loss (Eq. 7) represents the output of the hazard estimator corresponding to time $t$ for sample $i$, with covariates $x_i$, treatment indicator $a_i$, and dosage $q_i$.\n\n- The hazard estimator $h$ consists of treatment-specific network heads, each designed as a varying coefficient network with RBF basis functions, as described in Section 4.2.\n\n- Implementation details for $h$ are described in Appendix E.1. In all main experiments, we implement each of $h$'s network heads using three hidden layers, each comprising 100 nodes, which are modeled with five Gaussian RBFs that have learnable centers and shapes.\n\n### **Q6: Bias Variation Experiments**\n\n- Yes, this is precisely what we do: Figure 2 presents results for a fixed $\\xi=1.5$ (treatment bias) and varying $\\eta=0,1,2,3,4$ (dosage bias). The results confirm that DoseSurv outperforms the baselines across all levels of dosage assignment bias.\n\n### **Q7: Confidence Intervals**\n\n- Our confidence intervals are computed using standard statistical methods: we calculate the sample standard deviation (with Bessel's correction) and apply the appropriate t-distribution with degrees of freedom = $n-1$ for a small number of runs (n ≤ 30). The margin of error is computed as $t * (\\text{std} / \\sqrt{n})$, where $t$ is the critical value from the t-distribution for 95\\% confidence. While $n=5$ runs may seem limited, this is standard practice in deep learning research (see, e.g., [1]) due to computational constraints, and the t-distribution appropriately accounts for small sample uncertainty. \n\n\n## **Replies to Other Comments:**\n### **Treatment Setup Background**\n- Studying the effects of continuous treatments in healthcare is an important problem, yet progress is somewhat limited by the lack of robust models capable of addressing biases in observational data. Relevant areas include chemotherapy [2], insulin regimens [3], and cardiovascular treatments [4].\n\n- We provide a comprehensive literature review on HTE models (e.g., DRNet [5], SCIGAN [6], VCNet [7]), which involve treatment arms and dosage levels, in Section D of the Appendix. Additionally, we present an overview of selected key neural methods for HTE estimation under continuous-valued treatments or TTE outcomes in Table D.1.\n\n- Our framework addresses a critical gap in the literature between HTE estimation with continuous treatments and survival analysis.\n\n\n### **Causal Diagram**\n\n- The causal diagram is generally similar to those defined in related work on HTE estimation with survival data (e.g., [1,8]). However, the intervention now comprises both treatments $A$ and dosages $Q$. Following the notation in our paper, one can define the counting processes $N(t)=𝟙_{\\{\\tilde{T}\\leq t, \\Omega=1\\}}$, which tracks events (deaths), and $J(t)=𝟙_{\\{\\tilde{T}\\leq t, \\Omega=0\\}}$, which tracks censoring events. The general directed acyclic graph (DAG) is then given by:\n\n    ```\n                       ┌─────────┬─────────────┐\n                       │    ┌────▼───┐    ┌────▼───┐ \n                       │  ⋯ │ N(t-1) ├───▶│  N(t)  │ ⋯\n                       │    └─▲──┬───┘    └▲──▲──┬─┘\n        ┌────┐  ┌───┐  │      │  │         │  │  │\n        │ X₁─┼──┼▸A ├──┘      │  │  ┌──────┘  │  │\n        │ X₂─┼──┼▸Q ├──┐      │  │  │         │  │\n      X │    │  └───┘  │      │  │  │         │  │\n        │ X₃─┼─────────│──────┴──│──│─────────┘  │\n        │ X₄─┼─────────│──────┬──│──│─────────┐  │\n        └────┘         │      │  │  │         │  │\n                       │    ┌─▼──▼──┴┐    ┌───▼──▼─┐ \n                       │  ⋯ │ J(t-1) ├───▶│  J(t)  │ ⋯\n                       │    └────▲───┘    └─────▲──┘\n                       └─────────┴──────────────┘\n    \n  Causal Paths:\n\n    X₁ → A                  A, Q → N(t), J(t)  \n    X₂ → Q                  N(t) → J(t)\n    X₃ → N(t)               N(t-1) → N(t) \n    X₄ → J(t)               J(t-1) → N(t), J(t)\n\n    for t = 1, 2, ..., tₘ.\n\n    ```\n   Note that $X$ comprises (potentially overlapping) subgroups $X_1$, $X_2$, $X_3$, $X_4$ of covariates, which influence (1) the treatment assignment, (2) the dosage assignment, (3) the event process, and (4) the censoring process. These subgroups are associated with covariate shifts due to (1) treatment-selection bias, (2) dosage-selection bias, (3) event-induced bias, and (4) censoring bias, respectively [1]. The assumed causal graph is in line with Assumptions 1-6 in the paper. We are happy to include a proper illustration of the causal diagram in the final version of our paper.\n\n### **Data Description**\n\n- As one of our contributions is to adopt the Twins dataset as a benchmark dataset for HTE estimation with continuous treatments and survival data, we decided to include some descriptions in the main text and provide full details in the Appendix. However, we are open to moving the dataset descriptions to the Appendix in the final version.\n\n\n**References:**\n\n[1] A. Curth et al., \"SurvITE: Learning heterogeneous treatment effects from time-to-event data,\" in Adv. Neural Inf. Process. Syst., vol. 34, pp. 26740–26753, 2021.\n\n[2] L. R. Duska et al., \"A Phase I Study of Continuous Infusion Doxorubicin and Paclitaxel Chemotherapy with Granulocyte Colony-Stimulating Factor for Relapsed Epithelial Ovarian Cancer,\" Clin. Cancer Res., vol. 5, no. 6, pp. 1299–1305, Jun. 1999.\n\n[3] U. Kampmann *et al., \"Insulin dose–response studies in severely insulin-resistant type 2 diabetes—evidence for effectiveness of very high insulin doses,\" Diabetes Obes. Metab., vol. 13, no. 6, pp. 511–516, Apr. 2011.\n\n[4] M. Fiuzat *et al., \"Relationship of beta-blocker dose with outcomes in ambulatory heart failure patients with systolic dysfunction,\" J. Am. Coll. Cardiol., vol. 60, no. 3, pp. 208–215, Jul. 2012.\n\n[5] P. Schwab et al., \"Learning counterfactual representations for estimating individual dose-response curves,\" in Proc. AAAI Conf. Artif. Intell., vol. 34, no. 4, pp. 4046–4053, Feb. 2020.\n\n[6] I. Bica et al., \"Estimating the effects of continuous-valued interventions using generative adversarial networks,\" in Adv. Neural Inf. Process. Syst., vol. 33, pp. 16434–16445, 2020.\n\n[7] L. Nie et al., \"Varying coefficient neural network with functional targeted regularization for estimating continuous treatment effects,\" in Proc. ICLR., 2021.\n\n[8] O. M. Stitelman and M. J. van der Laan, \"Collaborative targeted maximum likelihood for time to event data\", Int. J. Biostat., vol. 6, no. 1, p. 12, Jan. 2010."}]

Author rebuttal:

Meta-review: **Meta-review for Submission 20419**

**Recommendation:** Need Discussion

### (a) Summary of the Paper

This paper introduces DoseSurv, a novel deep learning framework for estimating heterogeneous treatment effects (HTEs) for survival (time-to-event) outcomes under continuous-valued treatments, such as varying medication dosages. This addresses a critical but underexplored intersection in causal machine learning with significant applications in precision medicine. The proposed model, DoseSurv, uses a varying coefficient network with Radial Basis Functions (RBFs) to flexibly model the continuous dose-response relationship. To handle confounding bias from observational data, it incorporates a representation learning module regularized by an Integral Probability Metric (IPM) to balance covariate distributions across different treatments, dosages, and time points. The authors demonstrate DoseSurv's superior performance over strong baselines on both synthetic and semi-synthetic real-world data adapted from the Twins dataset.

### (b) The Case for Acceptance (The Prevailing View)

There is a strong and consistent positive sentiment across all four reviews, with final scores of 5, 4, 4, and 4. The consensus is that the paper is well-written, methodologically sound, and tackles a timely, important, and challenging problem. The reviewers praised the clarity of the presentation, the novelty of the approach (specifically the combination of the RBF-based hazard estimator and the IPM balancing for this setting), and the strength of the chosen baselines.

Crucially, the authors provided an exemplary and highly effective rebuttal that addressed nearly all reviewer concerns with new experiments and detailed clarifications. Key points that were successfully resolved include:
*   **Ablation and Justification of Design Choices:** In response to Reviewers Lefa and Bytv, the authors ran several new experiments, including ablations of the RBF estimator alone, an RBF T-learner baseline, a version with an MLP head, and a low-sample size experiment. These results successfully justified their core design choices.
*   **Clarification of Causal Assumptions and Methodology:** The authors provided a clear causal diagram (per Reviewer 3uNJ's request) and detailed explanations for their IPM-based balancing scheme versus alternatives like Generalized Propensity Scores (per Reviewer Wk85's request), successfully arguing for the suitability of their approach in a time-varying survival context.
*   **Handling of Censoring:** Concerns about the loss function's handling of censoring were clarified, with the authors explaining the standard use of the at-risk set and providing results from an experiment with reweighted loss terms.

The thoroughness of the rebuttal was highly convincing, leading Reviewer Wk85 to raise their score to a 5 (Accept), and Reviewers 3uNJ and Lefa to indicate they would also raise their scores. This signals a strong consensus that the paper, especially with the proposed revisions, is a high-quality contribution.

### (c) Points for Deeper Consideration

Despite the strong positive consensus, a "Need Discussion" is warranted to deliberate on the nuances of the paper's contribution and positioning, which were surfaced during the review process. These are not arguments for rejection but rather key aspects to consider when evaluating the paper's final placement and impact.

1.  **Justification of Design Choices in the Original Manuscript:** A recurring theme, particularly from Reviewer Bytv, was that the rationale behind several key design choices (e.g., the specific survival model architecture, the choice of IPM over reweighting) was not sufficiently explained in the initial submission. The authors provided excellent justifications in the rebuttal, which satisfied the reviewers. The point for discussion is whether a paper that requires such extensive post-submission clarification to fully motivate its methodology meets the bar for clarity at submission, or if this is a standard and acceptable outcome of the review process. The reviewers' positive final reactions suggest the latter, but it is a point worth noting.

2.  **Novelty as a Synthesis of Existing Ideas:** As clarified in the rebuttal (in response to Reviewer Lefa), the paper's novelty lies in skillfully adapting and combining powerful existing techniques—varying coefficient networks, RBFs, and IPM-based representation learning—to tackle the novel problem setting of continuous treatments in survival analysis. This is a strong and valuable form of contribution, but it is a synthesis rather than the invention of a fundamentally new algorithm. This is a standard profile for many successful papers but is important context for the ACs.

3.  **Practical Impact and Real-World Data Limitations:** Multiple reviewers noted the limited real-world validation. The authors correctly pointed out that this is a field-wide limitation due to the extreme scarcity of public datasets with continuous dosage, survival outcomes, and ground-truth counterfactuals. Their use and adaptation of the Twins dataset is already more rigorous than what is seen in much of the related literature. The discussion point here is how to position the paper's impact: it is a significant *methodological* contribution with high potential for clinical application, but it is not in itself a direct clinical validation study.

### (d) Key Points for AC Discussion

1.  **Final Recommendation Strength:** Given the unanimous positive feedback post-rebuttal and the authors' commitment to integrating their extensive clarifications, is this a clear accept? Do the remaining nuances (e.g., the need for the rebuttal to fully flesh out the design rationale) affect its standing for a potential spotlight or oral presentation?

2.  **Evaluating the Contribution Type:** The paper's strength is in identifying an important, overlooked problem and proposing the first robust, well-validated deep learning solution by synthesizing and extending existing concepts. How does this type of contribution compare to papers that propose a more fundamentally novel algorithm for an already well-established problem?

3.  **Framing the Limitations:** The limited real-world validation is a product of the field, not a failure of the authors. How should this be reflected in the final decision? It seems the consensus is to view this as an acceptable limitation for a paper focused on methodological advancement.
Decision: Accept (poster)