SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jonas Schweisthal, Dennis Frauen, Valentyn Melnychuk & Stefan Feuerriegel
Munich Center for Machine Learning
{jonas.schweisthal,fruamen,melnychuk,feuerriegel}@lmu.de
Decision-making in personalized medicine such as cancer therapy or critical care must often make choices for dosage combinations, i.e., multiple continuous treatments. Existing work for this task has modeled the effect of multiple treatments independently, while estimating the joint effect has received little attention but comes with non-trivial challenges. In this paper, we propose a novel method for reliable off-policy learning for dosage combinations. Our method proceeds along three steps: (1) We develop a tailored neural network that estimates the individualized dose-response function while accounting for the joint effect of multiple dependent dosages. (2) We estimate the generalized propensity score using conditional normalizing flows in order to detect regions with limited overlap in the shared covariate-treatment space. (3) We present a gradient-based learning algorithm to find the optimal, individualized dosage combinations. Here, we ensure reliable estimation of the policy value by avoiding regions with limited overlap. We finally perform an extensive evaluation of our method to show its effectiveness. To the best of our knowledge, ours is the first work to provide a method for reliable off-policy learning for optimal dosage combinations.
In personalized medicine, decision-making frequently involves complex choices for dosage combinations. For example, in cancer therapy, patients with chemotherapy or immunotherapy receive a combination of two or three drugs, and, for each, the dosage must be carefully personalized [16, 82]. In critical care, medical professionals must simultaneously control multiple parameters for mechanical ventilation, including respiratory rate and tidal volume [20].
One way to find optimal dosages (i.e., multiple continuous treatments) is through randomized controlled trials. Yet, such randomized controlled trials are often impractical due to the complexity of dosing problems or even unethical [62]. Instead, there is a direct need for learning optimal, individualized dosage combinations from observational data (so-called off-policy learning).
Several methods for off-policy learning aim at discrete treatments [1, 3, 4, 14, 23, 37, 40, 74, 76] and are thus not applicable to dosage combinations. In contrast to that, there are only a few methods that aim at off-policy learning for continuous treatments, that is, dosages [6, 10, 41, 44, 87]. To achieve this, some of these methods utilize the individualized dose-response function, which is typically not available and must be estimated from observational data. However, existing works model the effect of multiple continuous treatments independently. Because of that, existing works ignore the dependence structure among dosage combinations, that is, drug-drug interactions. For example, in chemotherapy, the combination of antineoplastic agents with drugs such as warfarin can have adverse effects on patient health [61]. In contrast to that, methods for off-policy learning that account for the joint effect of dosage combinations have received little attention and present the focus of our paper.
Crucially, off-policy learning for dosage combinations comes with non-trivial challenges. (1) Empirically, datasets in medical practice are often high-dimensional, yet of limited sample size. Hence, in practice, there are often sparse regions in the covariate-treatment space, which, in turn, may lead to estimated individualized dose-response functions that have regions with high uncertainty. (2) Mathematically, off-policy learning commonly relies upon the overlap assumption [65]. Yet, certain drug combinations are rarely used in medical practice and should be avoided for safety reasons. Both challenges lead to areas with limited observed overlap in the data, which results in unreliable estimates of the policy value. Motivated by this, we thus develop a method for reliable off-policy learning.
Proposed method: In this paper, we propose a novel method for reliable off-policy learning for dosage combinations. Our method proceeds along three steps. (1) We develop a tailored neural network that estimates the individualized dose-response function while accounting for the joint effect of multiple dependent dosages. Later, we use the network as a plug-in estimator for the policy value. (2) We estimate the generalized propensity score through conditional normalizing flows in order to detect regions with limited overlap in the shared covariate-treatment space. (3) We find the optimal individualized dosage combinations. To achieve that, we present a gradient-descent-ascent-based learning algorithm that solves the underlying constrained optimization problem, where we avoid regions with limited overlap and thus ensure a reliable estimation of the policy value. For the latter step, we use an additional policy neural network to predict the optimal policy for reasons of scalability.
Contributions: (1) We propose a novel method for reliable off-policy learning for dosage combinations. To the best of our knowledge, ours is the first method that accounts for the challenges inherent to dosage combinations. (2) As part of our method, we develop a novel neural network for modeling the individualized dose-response function while accounting for the joint effect of dosage combinations. Thereby, we extend previous works which are limited to mutually exclusive or independent treatments. (3) We perform extensive experiments based on real-world medical data to show the effectiveness of our method.
Estimating the individualized dose-response function: There are different methods to estimate the outcomes of continuous treatments, that is the individualized dose-response function. One stream uses the generalized propensity score [26, 28, 29], while others adapt generative adversarial networks [5]. An even other stream builds upon the idea of shared representation learning from binary treatments (e.g., [33, 53, 68, 69, 84]) and extends it to continuous treatments [57, 67]. However, all of the previous methods model the marginal effects of multiple dosages independently and thus without considering the dependence among when multiple treatments are applied simultaneously. Yet, this is unrealistic in medical practice as there are generally drug-drug interactions. To address this, we later model the joint effect of multiple dependent dosages.
There are some approaches for estimating the individualized effects of multi- or high-dimensional treatments, yet not for dosages. For example, some methods target conditional multi-cause treatment effect estimation [55, 58, 59, 73, 78, 86] but are restricted to discrete treatments and can not be directly applied to dosages. Other works use techniques tailored for high-dimensional structured treatments [21, 36, 54]. However, they do not focus on smooth and continuous estimates, or reliable off-policy learning, which is unlike our method.
Off-policy learning for dosages: There is extensive work on off-policy learning for discrete treatments [1, 3, 4, 14, 17, 23, 37, 40, 49, 50, 72, 74, 85]. While, in principle, these methods could be applied to our setting by discretizing the dosages, crucial information about the exact individually prescribed dosage values would be lost. This results in less accurate and less personalized dosage recommendations. However, there are only a few papers that address off-policy learning directly for continuous treatments ("dosages").
For dosages, previous works have derived efficient estimators for policy evaluation and subsequent optimization (e.g., [10, 41]). Other works focus directly on the policy optimization objective to find an optimal dosage [6, 44, 87]. In principle, some of these methods are applicable to settings with multiple dosages where they act as naive baselines. However, all of them have crucial shortcomings
$^{1}$ Code is available at https://github.com/JSchweinsthal/ReliableDosageCombi.
in that they do not account for multiple dependent dosages and/or that they do not aim at reliable decision-making under limited overlap. Both of the latter are contributions of our method.
Dealing with limited overlap: Various methods have been proposed to address limited overlap for estimating average treatment effects of binary treatments by making use of the propensity score (e.g., [7, 9, 24, 30, 38, 66]). In contrast, there exist only a few works for off-policy learning with limited overlap, typically for binary treatments [39, 42, 70].
Some works deal with the violation of the overlap assumption by leveraging uncertainty estimates for binary [31] and continuous treatment effect estimation [32]. However, these works do not tackle the problem of direct policy learning. In particular for dosage combinations, this is a clear shortcoming, as determining and evaluating policies becomes computationally expensive and infeasible with a growing number of samples and treatment dimensions.
Research gap: Personalized decision-making for optimal dosage combinations must (1) estimate the joint effect of dosage combinations, (2) then perform off-policy learning for such dosage combinations, and (3) be reliable even for limited overlap. However, to the best of our knowledge, there is no existing method for this composite task, and we thus present the first method aimed at reliable off-policy learning for dosage combinations.
Setting: We consider an observational dataset for  $n$  patients with i.i.d. observations  $\{(x_{i}, t_{i}, y_{i})\}_{i=1}^{n}$  sampled from a population  $(X, T, Y)$ , where  $X \in \mathcal{X} \subseteq \mathbb{R}^{d}$  are the patients' covariates,  $T \in \mathcal{T} \subseteq \mathbb{R}^{p}$  is the assigned  $p$ -dimensional dosage combination (i.e., multiple continuous treatments), and  $Y \in \mathbb{R}$  is the observed outcome. For example, in a mechanical ventilation setting,  $X$  are patient features and measured biosignals, such as age, sex, respiratory measurements, and cardiac measurements;  $T$  are the dosages of the ventilation parameters, such as respiratory rate and tidal volume; and  $Y$  is patient survival. For simplicity, we also refer to the dosages as treatments.
We adopt the potential outcomes framework [64], where  $Y(t)$  is the potential outcome under treatment  $T = t$ . In medicine, one is interested in learning a policy  $\pi : \mathcal{X} \to \mathcal{T}$  that maps covariates to suggested dosage combinations. Then, the policy value  $V(\pi) = \mathbb{E}[Y(\pi(X))]$  is the expected outcome of policy  $\pi$ . Our objective is to find the optimal policy  $\pi^*$  in the policy class  $\Pi$  that maximizes  $V(\pi)$ , i.e.,
$$
\pi^ {*} \in \underset {\pi \in \Pi} {\arg \max } V (\pi). \tag {1}
$$
In other words, we want to find  $\pi^{*}$  that minimizes the regret given by  $R(\pi^{*}) = \max_{\pi \in \Omega} V(\pi) - V(\pi^{*})$ , where  $\Omega$  denotes the set of all possible policies with  $\Pi \subseteq \Omega$ .
We further introduce the following notation, which is relevant to our method later. We denote the conditional potential outcome function by  $\mu(t,x) = \mathbb{E}[Y(t)\mid X = x]$ . For continuous treatments, this is also called individualized dose-response function. Further, the generalized propensity score (GPS) is given by  $f(t,x) = f_{T\mid X = x}(t)$ , where  $f_{T\mid X = x}$  denotes the conditional density of  $T$  given  $X = x$  with respect to the Lebesgue measure. The GPS is the stochastic policy used for assigning the observed treatments (i.e., logging / behavioral policy). In our method, both nuisance functions,  $\mu(t,x)$  and  $f(t,x)$ , are not known and must be later estimated from observational data.
Assumptions: To ensure identifiability of the outcome estimation, the following three standard assumptions are commonly made in causal inference [65]: (1) Consistency:  $Y = Y(T)$ . Consistency requires that the patient's observed outcome is equal to the potential outcome under the observed treatment. (2) Ignorability:  $Y(t) \perp T \mid X, \forall t \in \mathcal{T}$ . Ignorability ensures that there are no unobserved confounders. (3) Overlap:  $f(t,x) > \varepsilon, \forall x \in \mathcal{X}, t \in \mathcal{T}$  and for some constant  $\varepsilon \in [0,\infty)$ . Overlap (also known as common support or positivity) is necessary to ensure that the outcomes for all potential treatments can be estimated accurately for all individuals.
We distinguish between weak overlap for  $\varepsilon = 0$  and strong overlap for  $\varepsilon > 0$ . Since the estimation of the individualized dose-response function becomes more reliable with increasing  $\varepsilon$  in the finite sample setting, we refer to strong overlap in the following when we mention overlap, unless explicitly stated otherwise.
Challenges due to limited overlap: Limited overlap introduces both empirical and theoretical challenges. (1) Datasets with high dimensions and/or small sample size are especially prone to have
sparse regions in the covariate-treatment space  $\mathcal{X} \times \mathcal{T}$  and thus limited overlap [9]. As a result, the estimated individualized dose-response function will have regions with high uncertainty, which preclude reliable decision-making. (2) It is likely that  $\varepsilon$  is small for some regions as some drug combinations are rarely prescribed and should be avoided. To address this, we aim at an off-policy method that is reliable, as introduced in Sec. 4.
4 Reliable Off-Policy Learning for Dosage Combinations
In the following, we propose a novel method for reliable off-policy learning for dosage combinations. Our method proceeds along three steps (see Fig. 1): (1) We develop a tailored neural network that estimates the individualized dose-response function while accounting for the joint effect of multiple dependent dosages (Sec. 4.1). Later, we use the network as a plug-in estimator for the policy value. (2) We estimate the generalized propensity score through conditional normalizing flows in order to detect regions with limited overlap in the shared covariate-treatment space (Sec. 4.2). (3) We find the optimal individualized dosage combinations (Sec. 4.3). To achieve that, we present a gradient-descent-ascent-based learning algorithm that solves the underlying constrained optimization problem (Sec. 4.4), where we avoid regions with limited overlap and thus ensure a reliable estimation of the policy value. For the latter step, we use an additional policy neural network to predict the optimal policy for fast inference during deployment.
Figure 1: Overview of our method along steps (1)-(3).
4.1 Neural network for estimating the joint effect of multiple dependent dosages
In the first step, we aim to estimate the individualized dose-response function  $\mu(t,x)$ , which is later used as a plug-in estimator for the policy value  $V(\pi)$ . However, using a naive neural network with  $t$  and  $x$  concatenated as inputs can lead to diminishing effect estimates of  $t$  when  $x$  is high-dimensional [68]. In addition, in our setting with multiple continuous treatments, we want to yield smooth estimates of the dose-response surface to ensure better generalization when interpolating between sparse observations. To address this, we thus develop a tailored dosage combination network (DCNet) in the following.
Architecture of DCNet: Our DCNet builds upon the varying coefficient network (VCNet) [56] but extends it to learn the joint effect of multiple dependent dosages. DCNet consists of (i) a representation network  $\phi(x)$  and (ii) a prediction head  $h_{\eta(t)}(\phi)$ , where  $h$  is a network with  $d_{\eta}$  different parameters  $\eta(t) = (\eta_1(t), \dots, \eta_{d_{\eta}}(t))^T$ . The parameters  $\eta(t)$  are not fixed but depend on the dosage  $t$ . This is crucial to ensure that the influence of  $t$  does not diminish in a setting with high-dimensional covariates  $x$ .
To handle multiple dosages  $(t^{(1)},\ldots ,t^{(p)})$ , a naive way [2] would be to incorporate a separate prediction head for each treatment with corresponding parameters  $\eta^{(i)}(t^{(i)})$ . However, this would be ineffective for our task, as it captures only the marginal effect of a dosage but not the joint effect of the dosage combination. Instead, in order to learn the joint effect, we must not only model the marginal effect of a dosage (as in [2]) but also drug-drug interactions. We achieve that through a custom prediction head.
Prediction head in DCNet: In DCNet, we use a tailored prediction head, which is designed to model the joint effect of dosage combinations. For this, we leverage a tensor product basis (see, e.g., [48, 81]) to estimate smooth interaction effects in the dose-response surface. As a result, we can directly incorporate all  $p$  dosages simultaneously into the parametrization of the network  $h$  by using only one head containing the joint dosage information. This is unlike other networks that use  $p$  different independent heads [2] and thus cannot capture the joint effect but only marginal effects.
We define each scalar parameter  $\eta_j(t)$  of our prediction head as
$$
\eta_ {j} (t) = \sum_ {k _ {1} = 1} ^ {K _ {1}} \sum_ {k _ {2} = 1} ^ {K _ {2}} \dots \sum_ {k _ {p} = 1} ^ {K _ {p}} \beta_ {j, k _ {1} k _ {2} \dots k _ {p}} \cdot \psi_ {k _ {1}} ^ {(1)} (t ^ {(1)}) \cdot \psi_ {k _ {2}} ^ {(2)} (t ^ {(2)}) \cdot \dots \cdot \psi_ {k _ {p}} ^ {(p)} (t ^ {(p)}), \tag {2}
$$
where  $K_{1},\ldots ,K_{p}$  are the number of basis functions  $\psi_{k_i}^{(i)}$  of the  $p$  dosages  $t^{(1)},\dots,t^{(p)}$ , and  $\beta_{j,\dots}$  are the trainable coefficients for each element of the tensor product. In order to preserve the continuity of the non-linear dose-response surface, we choose  $\psi_{k_i}^{(i)}$  to be polynomial spline basis functions. Hence, we yield the parameter vector  $\eta (t) = \mathbf{B}\Psi (t)$ , where
$$
\mathbf {B} = \left[ \begin{array}{c c c} \beta_ {1, 1 \dots 1} & \dots & \beta_ {1, K _ {1} \dots K _ {p}} \\ \vdots & \ddots & \vdots \\ \beta_ {d _ {\eta}, 1 \dots 1} & \dots & \beta_ {d _ {\eta}, K _ {1} \dots K _ {p}} \end{array} \right] \in \mathbb {R} ^ {d _ {\eta} \times \left(\prod_ {i = 1} ^ {p} K _ {i}\right)} \tag {3}
$$
with matrix rows arranged as  $(\beta_{j,11\dots 1},\dots,\beta_{j,K_11\dots 1},\dots,\beta_{j,K_1\dots K_{p - 1}1},\dots,\beta_{j,K_1\dots K_{p - 1}K_p})$  with
$$
\Psi (t) = \left(\psi^ {(1)} \left(t ^ {(1)}\right) \otimes \psi^ {(2)} \left(t ^ {(2)}\right) \otimes \dots \otimes \psi^ {(p)} \left(t ^ {(p)}\right)\right) \in \mathbb {R} ^ {\left(\prod_ {i = 1} ^ {p} K _ {i}\right)} \tag {4}
$$
and with  $\psi^{(i)} = (\psi_{1}^{(i)},\dots ,\psi_{K_{i}}^{(i)})^{T}$ , where  $\otimes$  is the Kronecker product. For training DCNet, we minimize the mean squared error loss  $\mathcal{L}_{\mu} = \frac{1}{n}\sum_{i = 1}^{n}\left(\hat{\mu} (t_i,x_i) - y_i\right)^2$ .
Our prediction head has three key benefits: (1) When estimating  $\mu(t, x)$ , the tensor product basis enables us to model non-linear interaction effects among different dosage dimensions while still having a dose-response surface that offers continuity, smoothness, and expressiveness. (2) Having a dose-response surface that is smooth between different dosages is beneficial for our task as it allows for better interpolation between sparse data points. (3) Smoothness is also exploited by our gradient-based learning algorithm later. The reason is that a smooth plug-in estimator  $\hat{\mu}(t, x)$  for the estimated policy value  $\hat{V}(\pi)$  allows us to have more stable gradients and thus more stable parameter updates during learning.
4.2 Conditional normalizing flows for detecting limited overlap
In the second step, we estimate the GPS  $f(t,x)$  in order to detect limited overlap, which later allows us to avoid regions with high uncertainty and thus to yield a reliable policy.
Why is the GPS relevant for reliable decision-making? Due to sparsity in the data or otherwise limited overlap, the estimates of the individualized dose-response function  $\hat{\mu}(t,x)$  from our DCNet may be naturally characterized by regions with high uncertainty, that is, regions in the covariate-treatment space with low GPS. If we would use the estimated  $\hat{\mu}(t,x)$  in a naïve manner as a plug-in estimator for policy learning, we may yield a policy where there can be a large discrepancy between the predicted outcome  $\hat{\mu}(t,x)$  and the true outcome  $\mu(t,x)$  and, therefore, a large discrepancy between  $\hat{V}(\pi)$  and  $V(\pi)$ . By avoiding regions with a low GPS, we can thus ensure reliable decision-making in safety-critical applications such as medicine.
Estimation of GPS: The GPS (i.e., the logging policy) is typically not known in real-world settings. Instead, we must estimate the GPS from observational data. To this end, we leverage conditional
normalizing flows (CNFs) [75, 80]. CNFs are a fully-parametric generative model built on top of normalizing flows [60, 71] that can model conditional densities  $p(y \mid x)$  by transforming a simple base density  $p(z)$  through an invertible transformation with parameters  $\gamma(x)$  that depend on the input  $x$ .<sup>2</sup>
In our method, we use CNFs with parameters  $\gamma(x)$  to estimate the probability density function  $f(t,x) = f_{T|X=x}(t)$  of  $T$  conditioned on the patient covariates  $X = x$ . For our method, CNFs have several advantages. First, CNFs are universal density approximators [11, 12, 15, 27], making them flexible and able to capture even complex density functions for dosage combinations. Second, CNFs are properly normalized. Third, unlike kernel methods, CNFs are fully parametric, and, once trained, the inference time is constant [52]. This is particularly important later in the third step of our method (Sec. 4.3): For each training step and training sample in the batch, the predicted conditional density has to be evaluated again under the updated policy, and, hence, a fast inference time of density is crucial to achieving a significant speed up when training the policy network.
In our method, we use neural spline flows [15], as they allow for flexible estimation of multimodal densities, and, in combination with masked auto-regressive networks [13], we can extend them to multiple dimensions. This combination allows for fast and exact evaluation of the multidimensional densities. We also set the base density,  $p(z)$ , to standard normal distribution. The CNFs are trained by minimizing the negative log-likelihood loss  $\mathcal{L}_f = -\frac{1}{n}\sum_{i=1}^n\log\hat{f}(t_i,x_i)$ .
In the following section, we use the estimated GPS  $\hat{f}(t,x)$  during policy learning in order to avoid regions with limited overlap and thus ensure reliable decision-making.
4.3 Constrained optimization to avoid regions with limited overlap
Reliable decision-making: Our aim is to find the policy  $\pi^{\mathrm{rel}}$  which maximizes the estimated policy value  $\hat{V}(\pi)$  while ensuring that  $V(\pi)$  can be estimated reliably. For this, we suggest to constrain our search space in the shared covariate-treatment space to regions that are sufficiently supported by data, i.e., where overlap is not violated. Formally, we rewrite our objective from Eq. (1) as
$$
\pi^ {\text {r e l}} \in \underset {\pi \in \Pi^ {\mathrm {r}}} {\arg \max } \hat {V} (\pi) \quad \text {w i t h} \quad \Pi^ {\mathrm {r}} = \left\{\pi \in \Pi \mid f (\pi (x), x) > \bar {\varepsilon}\right), \forall x \in \mathcal {X} \}, \tag {5}
$$
where  $\overline{\varepsilon}$  is a reliability threshold controlling the minimum overlap required for our policy. In our setting with finite observational data, this yields the constrained optimization problem
$$
\max  _ {\pi} \frac {1}{n} \sum_ {i = 1} ^ {n} \hat {\mu} (\pi (x _ {i}), x _ {i}) \quad \text {s . t .} \quad \hat {f} (\pi (x _ {i}), x _ {i}) \geq \bar {\varepsilon}, \forall i \in \{1, \dots , n \}. \tag {6}
$$
Here, we use our DCNet as a plug-in dose-response estimator  $\hat{\mu}(t,x)$ , and we further use our conditional normalizing flows to restrict the search space to policies with a sufficiently large estimated GPS  $\hat{f}(t,x)$ .