SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Neural network for reliable policy learning: To learn the optimal reliable policy, we further suggest an efficient procedure based on gradient updates in the following. Specifically, we train an additional policy neural network to predict the optimal policy. Our choice has direct benefits in practice. Why not just optimize over the dosage combinations for each patient individually? In principle, one could also use methods for constrained optimization and leverage the nuisance functions  $\hat{\mu}$  and  $\hat{f}$  to directly optimize over the dosage assignments  $t$  for each patient, i.e.,  $t_i^{\mathrm{rel}} = \arg \max_{\tilde{t}_i}[\hat{\mu} (\tilde{t}_i,x_i)]$  s.t.  $\hat{f} (\tilde{t}_i,x_i)\geq \overline{\varepsilon}$ . However, this does not warrant scalability at inference time, as it requires that one solves the optimization problem for each incoming patient. The complexity of this optimization problem scales with both the number of dosages and the number of samples in the data to be evaluated. This makes treatment recommendations costly, which prohibits real-time
${}^{2}$  Normalizing flows were introduced for expressive variational approximations in variational autoencoders [60, 71]. We provide a background in Appendix A.
<sup>3</sup>Most prior works neglect a proper estimation of the GPS by mainly using non-flexible parametric models [26] or kernel density estimators [41]. Similarly, the varying coefficient network from [56] uses an additional density head but enforces a discretization. It is thus overly simplified and, therefore, not reliable and not continuous. Both are crucial shortcomings for our task. As a remedy, we propose the use of CNFs.
decision-making in critical care. As such, direct methods for optimization are rendered impractical for real-world medical applications. Instead, we suggest to directly learn a policy that predicts the optimal dosage combination for incoming patients.
As the constrained optimization problem in Eq. (6) cannot be natively learned by gradient updates, we must first transform Eq. (6) into an unconstrained Lagrangian problem
$$
\begin{array}{l} \min  _ {\theta} - \frac {1}{n} \sum_ {i = 1} ^ {n} \hat {\mu} \left(\pi_ {\theta} \left(x _ {i}\right), x _ {i}\right) \quad \text {s . t .} \hat {f} \left(\pi_ {\theta} \left(x _ {i}\right), x _ {i}\right) \geq \bar {\varepsilon}, \forall i \\ \Longleftrightarrow \min  _ {\theta} \max  _ {\lambda_ {i} \geq 0} - \frac {1}{n} \sum_ {i = 1} ^ {n} \left\{\hat {\mu} \left(\pi_ {\theta} \left(x _ {i}\right), x _ {i}\right) - \lambda_ {i} \left[ \hat {f} \left(\pi_ {\theta} \left(x _ {i}\right), x _ {i}\right) - \bar {\varepsilon} \right] \right\}, \tag {7} \\ \end{array}
$$
where  $\pi_{\theta}(x_i)$  is the policy learner with parameters  $\theta$ , and  $\lambda_{i}$  are the Lagrange multipliers for sample  $i$ . The Lagrangian min-max-objective can be solved by adversarial learning using gradient descent-ascent optimization (see, e.g., [47] for a background). Details are presented in the next section where we introduce our learning algorithm.
4.4 Learning algorithm
The learning algorithm for our method proceeds along the following steps (see Algorithm 1):
(1) We use our DCNet to estimate the individualized dose-response function  $\hat{\mu}(t_i, x_i)$  for dosage combinations.
(2) We estimate the GPS  $\hat{f} (t_i,x_i)$  using the CNFs to detect areas with limited overlap.
(3) We plug the estimated  $\hat{\mu}(t,x)$  and  $\hat{f}(t,x)$  into the Lagrangian objective from Eq. (7). As the policy learner, we use a multilayer-perceptron neural network  $\hat{\pi}_{\theta}(x)$  to predict the optimal dosage combination  $t$  given the patient covariates  $x$ . We use stochastic gradient descent-ascent. In detail, we perform adversarial learning with alternating parameter updates. First, in every training step, we perform a gradient descent step to update the parameters  $\theta$  of the policy network  $\hat{\pi}_{\theta}$  with respect to the policy loss
$$
\mathcal {L} _ {\pi} (\theta , \lambda) = - \frac {1}{n} \sum_ {i = 1} ^ {n} \left\{\hat {\mu} \left(\hat {\pi} _ {\theta} \left(x _ {i}\right), x _ {i}\right) - \lambda_ {i} \left[ \hat {f} \left(\hat {\pi} _ {\theta} \left(x _ {i}\right), x _ {i}\right) - \bar {\varepsilon} \right] \right\}, \tag {8}
$$
where  $\overline{\varepsilon}$  is the reliability threshold and where  $\lambda = \{\lambda_i\}_{i=1}^n$  are trainable penalty parameters. Then, we perform a gradient ascent step to update the penalty parameters  $\lambda$  with respect to  $\mathcal{L}_{\pi}$ . Here, each  $\lambda_i$  is optimized for each sample  $i$  in the training set.
As there is no guarantee for global convergence in the non-convex optimization problem, we perform  $k = 1,\dots ,K$  random restarts and select the best run as evaluated on the validation set by
$$
\pi_ {\theta} ^ {\mathrm {r e l}} = \pi_ {\theta} ^ {(j)}, \quad \text {w i t h} \quad j = \arg \max  _ {k} \sum_ {i = 1} ^ {n} \hat {\mu} \left(\pi_ {\theta} ^ {(k)} \left(x _ {i}\right), x _ {i}\right) \cdot \mathbb {1} \left\{\hat {f} \left(\pi_ {\theta} ^ {(k)} \left(x _ {i}\right), x _ {i}\right) \geq \bar {\varepsilon} \right\}, \tag {9}
$$
where  $\pi_{\theta}^{(k)}$  is the learned policy in run  $k$  and  $\mathbb{1}\{\cdot\}$  denotes the indicator function. As a result, we select the policy  $\pi_{\theta}^{(j)}$  which maximizes  $\hat{V}(\pi)$  under the overlap constraint  $\hat{f}\left(\pi_{\theta}^{(k)}(x_i), x_i\right) \geq \bar{\varepsilon}, \forall i$ . Note that we use the validation loss from Eq. (9) because  $\mathcal{L}_{\pi}$  cannot be directly applied for evaluating the performance on the validation set, as  $\lambda_i$  can only be optimized for observations in the training set during the policy learning.
We perform extensive experiments using semi-synthetic data from real-world medical settings to evaluate the effectiveness of our method. Semi-synthetic data are commonly used to evaluate causal inference methods as it ensures that the causal ground truth is available [8, 83] and thus allows us to benchmark the performance.
MIMIC-IV: MIMIC-IV [35] is a state-of-the-art dataset with de-identified health records from patients admitted to intensive care units. Analogous to [67], we aim to learn optimal configurations of mechanical ventilation to maximize patient survival  $(Y)$ . Here, treatments  $(T)$  involve the configuration of mechanical ventilation in terms of respiratory rate and tidal volume. When filtering for  $T$ , we yield a dataset with  $n = 5476$  patients. We select 33 variables as patient covariates  $X$  (e.g., age, sex, respiratory measurements, cardiac measurements). Unlike [67] where treatments  $T$  are modeled as  $p$  mutually exclusive dosages, we model the more realistic setting, where  $T$  are dosage combinations with a joint effect. Our choice is well aligned with medical literature according to which the parameters of mechanical ventilation are dependent and must be simultaneously controlled [20]. Further details for MIMIC-IV are in Appendix B.
TCGA: TCGA [79] is a diverse collection of gene expression data from patients with different cancer types. Our aim is to learn the optimal dosage combination for chemotherapy  $(T)$ . In chemotherapy, combining different drugs is the leading clinical option for treating cancer [82]. Let patient survival denote the outcome  $(Y)$ . We select the same version of the dataset as in [5, 67], containing the measurement of 4,000 genes as features  $X$  from  $n = 9659$  patients. Unlike previous works [5, 67], we again model  $T$  as different, simultaneously assigned dosages with a joint effect. This follows medical practice where drug combinations in cancer therapy must be jointly optimized for each patient profile [16, 82]. Given that TCGA is high-dimensional, we expect it to be an especially challenging setting for evaluation. Further details for TCGA are in Appendix B.
In both datasets, the response  $Y \sim N(\mu(T, X), 0.5)$  can be interpreted as the log-odds ratio of patient survival, which is to be maximized. Due to the fact that we work with actual data from medical prac
tice, certain dosage combinations conditional on the patients' features are more likely to be observed. Even other dosage combinations may not be assigned at all as they can cause harm. For our experiments, we introduce a parameter  $\alpha$  controlling the dosage bias, which changes the observed policy of dosage assignments from random uniform  $(\alpha = 0)$  to completely deterministic  $(\alpha = \infty)$ . Hence, with increasing  $\alpha$ , the areas in the covariate-treatment space will be increasingly prone to limited overlap. We refer to Appendix B for details. When not stated otherwise, we choose  $\alpha = 2$  and  $p = 2$  as the default setting for our experiments.
Baselines: To the best of our knowledge, there are no existing methods tailored to off-policy learning for dosage combinations under limited overlap. Hence, we compare our method with the following baselines, which vary different steps in our method and thus allow us to attribute the source of performance gain. Specifically, we vary dose-response estimation  $(\hat{\mu}(t,x))$  vs. optimization  $(\hat{\pi})$ . (1) For estimating  $\hat{\mu}(t,x)$ , we use two baselines: (i) a simple multi-layer perceptron (MLP) as a standard neural alternative to our DCNet, and (ii) VCNet [56] as the SOTA method for dose-response estimation for independent treatments with a separate head per treatment dimension. (2) For optimization, we minimize the policy loss  $-\frac{1}{n}\sum_{i=1}^{n}\hat{\mu}(\hat{\pi}_{\theta}(x_i), x_i)$  without constraints for dealing with limited overlap (naive).
Evaluation: For training, we split the data into train / val / test sets (64 / 16 / 20%). We first train the nuisance models  $\hat{\mu}(t,x)$  and  $\hat{f}(t,x)$ . We then perform  $k = 5$  random restarts to train  $\hat{\pi}(t,x)$ . We select the best run with respect to the policy loss on the factual validation dataset. Thereby, we carefully adhere to the characteristics in real-world settings, where one does not have access to the true dose-response function  $\mu(t,x)$  on the validation set for model selection. In our method, we set
Algorithm 1: Reliable off-policy learning for dosage combinations
Input : data  $(X,T,Y)$  , reliability threshold  $\overline{\varepsilon}$
Output: optimal reliable policy  $\hat{\pi}_{\theta}^{\mathrm{rel}}$
// Step 1: Estimate individualized dose-response
function using our DCNet
Estimate  $\hat{\mu} (t,x)$  via loss  $\mathcal{L}_{\mu}$
// Step 2: Estimate GPS using conditional normalizing flows
Estimate  $\hat{f} (t,x)$  via loss  $\mathcal{L}_f$
// Step 3: Train policy network using our reliable learning algorithm
for  $k\in \{1,2,\ldots ,K\}$  do
// K runs with random initialization
$\hat{\pi}_{\tau}^{(k)} \gets$  initialize randomly
$\lambda^{\theta}\gets$  initialize randomly
for each batch do
$$
\begin{array}{l} \left| \quad \right. \quad \left. \quad \right| \quad \mathcal {L} _ {\pi} \leftarrow - \frac {1}{n} \sum_ {i = 1} ^ {n} \left\{\hat {\mu} \left(\hat {\pi} _ {\theta} ^ {(k)} \left(x _ {i}\right), x _ {i}\right) \right. \\ \left. - \lambda_ {i} \left[ \hat {f} \left(\hat {\pi} _ {\theta} ^ {(k)} \left(x _ {i}\right), x _ {i}\right) - \bar {\varepsilon} \right] \right\} \\ \left| \quad \theta \leftarrow \theta - \eta_ {\theta} \nabla_ {\theta} \mathcal {L} _ {\pi} \right. \\ \left| \quad \lambda \leftarrow \lambda + \eta_ {\lambda} \nabla_ {\lambda} \mathcal {L} _ {\pi} \right. \\ \end{array}
$$
// select best learned policy wrt constrained
objective on validation set
$\hat{\pi}_{\theta}^{\mathrm{rel}} \gets \pi_{\theta}^{(j)}$ , with  $j =$
$$
\arg \max  _ {k} \sum_ {i = 1} ^ {n} \hat {\mu} \left(\pi_ {\theta} ^ {(k)} \left(x _ {i}\right), x _ {i}\right) \cdot \mathbb {1} \left\{\hat {f} \left(\pi_ {\theta} ^ {(k)} \left(x _ {i}\right), x _ {i}\right) \geq \bar {\varepsilon} \right\}
$$
the reliability threshold  $\bar{\varepsilon}$  to the  $5\%$  -quantile of the estimated GPS  $\hat{f} (t,x)$  of the train set. Further details including hyperparameter tuning are in Appendix E.
For evaluation, we compare the methods on the test set using the regret  $R(\hat{\pi}) = \max_{\pi \in \Omega} V(\pi) - V(\hat{\pi})$ . We report: (1) the regret for the selected policy, which is the policy with the best policy loss on the validation set out of all  $k$  policies; (2) the average regret across all  $k$  restarts; and (3) the standard deviation to better assess the reliability of the method.
Results: Table 1 compares the performance of our method against the baselines. We make the following observations: (1) DCNet performs similarly to MLP for low-dimensional data such as MIMIC-IV, as expected. However, we find substantial improvements from DCNet over
MLP for high-dimensional data such as TCGA, showing the advantage of enforc
Methods	MIMIC-IV	TCGA
μ	π	Selected	Mean	Std	Selected	Mean	Std
Oracle (μ)	observed	1.21	-	-	1.06	-	-
MLP	naive	0.02	1.44	1.56	1.94	1.88	0.12
MLP	reliable	0.03	0.03	0.00	1.78	1.78	0.03
VCNet	naive	2.13	2.49	0.92	3.64	2.47	1.19
VCNet	reliable	0.07	0.07	0.00	0.06	0.06	0.00
DCNet	naive	2.81	2.07	1.17	2.52	0.94	1.21
DCNet	reliable (ours)	0.04	0.04	0.00	0.03	0.03	0.00
Reported: best / mean / standard deviation (lower = better)
Table 1: Performance against baselines. Regret on test set over  $k = 5$  restarts.
ing smoothness and expressiveness in our prediction head. (2) DCNet outperforms VCNet on both datasets, which demonstrates the importance of modeling the joint effect of dosage combinations. (3) Choosing a policy that is reliable greatly benefits overall performance. For example, the baselines (+naive) even perform worse than the observed policy for TCGA, whereas our method (+reliable) is clearly superior. (4) We also observe that our selection of reliable policies greatly reduces the standard deviation of the observed regrets. This further demonstrates the reliability of our proposed method. Overall, our method works best across both datasets. For example, the regret on the TCGA dataset for the selected policy drops from 2.52 (DCNet+naive) to only 0.03 (DCNet+reliable), which is a reduction by  $98.8\%$ . In sum, this demonstrates the effectiveness of our method.
Comparison to discretization: There exist several methods for off-policy learning for discrete treatments (see Sec. 2), including methods targeted for overlap violations and treatment combinations. In principle, we could also discretize the dosages in our setting and apply such methods. However, we would then lose information about the exact dosage values and could not exploit the continuity of the dose-response surface. To demonstrate the benefits of our method, we benchmark against an "oracle" discretized baseline where we assume perfect knowledge of the dose-response function and perfect individual treatment assignment.
This serves as an upper bound for all policy learning methods with discretization. We display the results for different granularities of equally-spaced grid discretization in Table 2. We observe that, even in the oracle setting, discretization itself leads to a higher regret than our proposed method, which confirms the advantages of leveraging continuous dosage information.
Methods	MIMIC-IV	TCGA
μ	π	Regret
Oracle (μ)	discrete (3x3)	3.32	2.94
Oracle (μ)	discrete (4x4)	0.96	0.95
Oracle (μ)	discrete (5x5)	0.05	0.06
DCNet	reliable (ours)	0.04	0.03
Reported: best / mean / standard deviation (lower = better)
Robustness checks: We evaluate the robustness of our method across settings with different levels of limited overlap. For this, we vary the dosage bias  $\alpha$
(in Fig. 2) and the number of dosages  $p$  (in Fig. 3). Here, we use DCNet for modeling  $\hat{\mu}(t,x)$  in both methods, but then vary the policy selection (naive vs. reliable, as defined above). Our method (DCNet+reliable) shows robust behavior. As desired, it achieves a low regret and a low variation in all settings. In contrast, the naive method leads to high variation. This must be solely attributed to that our constrained optimization for avoiding regions with limited overlap leads to large performance gains. We also observe another disadvantage of the naive method: it may select a suboptimal policy out of the different runs due to the fact that the estimation of  $\hat{\mu}(t,x)$  is erroneous in regions with limited overlap. Again, this underlines the benefits of our method for reliable off-policy learning of dosage combinations in settings with limited overlap. In Appendix C, we demonstrate the robustness of our method in further setups.
Table 2: Performance against methods based on discretized treatments.
Shown: selected policy (line) and the range over 5 runs (area). The naive baseline has a large variability across runs while ours is highly robust as expected (i.e., we see no variability).
(a) MIMIC-IV
(b) TCGA
Figure 2: Robustness for dosage bias  $\alpha$ .
(a) MIMIC-IV
Figure 3: Robustness for number of dosages  $p$ .
(b) TCGA
**Limitations:** Our method makes an important contribution over existing literature in personalized treatment design for dosage combinations by adjusting for the naturally occurring but non-trivial problems of drug-drug interactions and limited overlap. However, the complexity of real-world data in medical practice can limit the applicability of our method in certain ways. (i) Our method does not account for the ignorability assumption, which can result in biased estimates in the case of unobserved confounders or other missing data. (ii) We apply our method in a static treatment setting, whereas, in several clinical applications, time-series data are available, e.g., sequences of varying dosages, multiple treatment cycles, and right-censored data. Thus, future work could extend our method to account for causal effect estimates in a time-varying setting (e.g., [18, 25, 51]). (iii) Our method relies on the – for medical applications reasonable – implicit assumption of smooth dose-response surfaces. In settings with extremely unsmooth or even stepwise dose-response surfaces, other methods may be more suitable for dose-response estimation. (iv) Our method aims at reliable policy learning by avoiding areas with limited overlap to minimize potential harm. In other applications such as marketing, one may want to explicitly target unreliable regions to acquire new customers.
Broader impact: Our method can have a significant impact in the field of medicine, where reliable dosing recommendations are crucial for effective treatment. Our approach addresses the limitations of previous methods in terms of reliability and estimating joint effects of dosage combinations, and, hence, can improve decision support for clinicians. The reliability of dosage recommendations is particularly important, as incorrect dosing can have serious consequences for patients, such as adverse effects or treatment failure. By addressing this gap, our approach provides a valuable tool for decision support in medicine, empowering healthcare professionals to make more personalized and safe treatment recommendations.
[79] John N. Weinstein et al. “The Cancer Genome Atlas Pan-Cancer analysis project”. In: Nature Genetics 45.10 (2013), pp. 1113–1120.
[80] Christina Winkler et al. "Learning likelihoods with conditional normalizing flows". In: arXiv preprint arXiv:1912.00042 (2019).
[81] Simon N. Wood. "Low-rank scale-invariant tensor product smooths for generalized additive mixed models". In: Biometrics 62.4 (2006), pp. 1025-1036.
[82] Debra Wu et al. “Design principles of drug combinations for chemotherapy”. In: Journal of Controlled Release 323 (2020), pp. 36–46.
[83] Liyuan Xu et al. "Learning deep features in instrumental variable regression". In: ICLR. 2021.
[84] Yao Zhang, Alexis Bellot, and Mihaela van der Schaar. "Learning overlapping representations for the estimation of individualized treatment effects". In: AISTATS. 2020.
[85] Zhengyuan Zhou, Susan Athey, and Stefan Wager. "Offline Multi-Action Policy Learning: Generalization and Optimization". In: Operations Research 71.1 (2023), pp. 148-183.
[86] Hao Zou et al. "Counterfactual prediction for bundle treatment". In: NeurIPS. Vol. 33. 2020.
[87] Hao Zou et al. "Counterfactual prediction for outcome-oriented treatments". In: ICML. 2022.
Estimating the dose-response function under dosage bias: Estimating the dose-response function becomes challenging with increasing dosage bias. Although the unbiasedness of counterfactual estimation can be achieved through the minimization of the prediction loss on the factual data, the inherent scarcity of data in continuous treatment settings leads to increased variance in the estimation process (see, e.g., [2]), especially in regions with limited overlap. Consequently, the elevated variance hinders accurate generalization over the whole covariate-treatment space and contributes to a larger generalization error, which, in turn, impedes the reliable estimation of treatment effects. Effectively addressing this issue is crucial to improve the robustness and validity of counterfactual predictions.
In Sec. 2, we introduced various existing methods for estimating the dose-response function, e.g., [5, 26, 28, 29, 57, 67]. Besides a custom neural architecture, some methods use additional techniques to reduce the variance of counterfactual predictions for discrete or continuous treatments. Examples include reweighting [63, 77], importance sampling [22], targeted regularization [56], or balancing representations [2, 33, 67]. In our DCNet, we build on top of the neural architectures of [57, 67] and then develop a method that is capable of modeling the joint effect of dosage combinations.
In theory, we could also adapt additional techniques for variance reduction of the dose-response prediction to our setting of dosage combinations and leverage them in combination with DCNet. However, this may not be beneficial for our goal of reliable policy learning for dosage combinations under limited overlap, which is different from the goal of accurate dose-response estimation over the whole covariate-treatment space. These techniques aim to reduce the variance of the dose-response prediction in scarce regions (i.e., with limited overlap), which, as a consequence, shifts the focus away from regions with sufficient overlap. As we constrain the policy to regions with sufficient overlap anyway by using our reliable policy optimization, we aim for an adequate dose-response estimation in these regions, while regions with limited overlap can be neglected. Hence, in our method for reliable policy learning, techniques for tackling dosage bias during dose-response estimation can be considered more counterproductive than useful. We demonstrate this empirically in Appendix D.
Normalizing flows: Normalizing flows were introduced for expressive variational approximations in variational autoencoders [60, 71] and are based on the principle of transforming a simple base distribution, such as a Gaussian, through a series of invertible transformations to approximate the target data distribution. Let us denote the base distribution as  $p_z(z)$ , where  $z$  is a latent variable, and the target data distribution as  $p_x(x)$ , where  $x$  represents the observed data. The goal is to find a mapping  $f$  such that  $x = f(z)$ .
The core idea behind normalizing flows is to model the transformation  $f$  as a composition of invertible functions  $f = f_{K} \circ f_{K-1} \circ \ldots \circ f_{1}$ , where  $K$  is the number of transformations. The last transformation  $f_{K}$  is typically modeled as an affine transformation, and others are invertible nonlinear transformations. The inverse transformation can be easily computed since each  $f_{k}$  is invertible.
The change of variables formula allows us to compute the probability density of  $x$  in terms of  $z$ . If we assume the base distribution  $p_z(z)$  is known and has a simple form (e.g., a Gaussian), we can compute the density of  $x$  via
$$
p _ {x} (x) = p _ {z} (z) \left| \det  \left(\frac {\partial f}{\partial z}\right) \right| ^ {- 1} \tag {10}
$$