SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\left|\operatorname{det}\left(\frac{\partial f}{\partial z}\right)\right|$  is the determinant of the Jacobian matrix of the transformation  $f$  with respect to  $z$ . This determinant captures the stretching and squishing of the latent space induced by the transformation, thus allowing us to adjust the density of  $z$  to match the density of  $x$ .
To perform inference, such as density estimation or sampling, we need to compute the log-likelihood of the observed data  $x$ . Given a dataset  $\mathcal{D} = \{x_1, \ldots, x_n\}$ , the log-likelihood can be computed as the sum of the log-densities of each data point
$$
\log p (\mathcal {D}) = \sum_ {i = 1} ^ {n} \log p _ {x} \left(x _ {i}\right). \tag {11}
$$
Conditional normalizing flows: In our method, we leverage conditional normalizing flows (CNFs) [75, 80] to estimate the generalized propensity score (GPS). CNFs 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$ .
In our method, we use neural spline flows [15] in combination with masked auto-regressive networks [13]. Here, masked auto-regressive networks are used to model the conditional distribution of each variable in the data given the conditioning variable. The network takes as input the conditioning variable and generates each variable in the data sequentially, one at a time, while conditioning on the previously generated variables. This autoregressive property allows for efficient computation and generation of samples from the conditional distribution. For the motivation and advantages of our modeling choice, we refer to Sec. 4.2.
Offline reinforcement learning: A related literature stream covers off-policy evaluation (OPE) in an offline reinforcement learning (ORL) setting. Some works leverage value functions and/or logging policies to achieve reliable policy learning (e.g., [19, 45, 46]). The main differences to our setting are that ORL assumes a Markov decision process and sequential decision-making. In contrast, we focus on non-sequential, off-policy learning from observational data. Also, unlike ORL or standard OPE, we leverage the causal structure of treatment assignment and the dose-response function to return causal estimands. This allows us not only to learn decisions but also causal estimates for potential outcomes. This is crucial for medical practice, where physicians would like to reason over different treatments rather than blindly following a fully automated system.
B Data Description
MIMIC-IV: MIMIC-IV $^{4}$  is the latest version of the publicly available dataset of The Medical Information Mart for Intensive Care [35] and contains de-identified health records of patients admitted to intensive care units (ICUs). The dataset includes comprehensive clinical data, such as vital signs, laboratory measurements, medications, diagnoses, procedures, and more. In line with the approach in [67], our primary objective is to acquire knowledge about the optimal configurations of mechanical ventilation that can effectively maximize patient survival  $(Y)$ . Note that some works (e.g., [5, 67]) used the previous version of the MIMIC dataset (MIMIC-III [34]), whereas we use the most current one to offer realistic, comparable, and meaningful findings.
As treatments  $(T)$ , we select the specific settings (dosages) utilized for mechanical ventilation, encompassing parameters such as respiratory rate and tidal volume. By filtering the data for the last available measurements per patient before ventilation, we are able to compile a dataset consisting of  $n = 5476$  patients. We carefully selected 33 patient covariates  $(X)$  that encompass factors of medical relevance such as age, sex, and various respiratory and cardiac measurements. For detailed identifiers of the selected covariates, we refer to our code.
While [67] adopted a framework in which treatments  $T$  were modeled as a set of  $p$  mutually exclusive dosages, we opted for a more realistic approach, considering that the configuration of mechanical ventilation involves dosage combinations with a joint effect. This decision aligns closely with existing medical literature, which emphasizes the interdependence of mechanical ventilation parameters and underscores the importance of simultaneous control [20].
TCGA: The Cancer Genome Atlas (TCGA) dataset[79] represents a comprehensive and diverse collection of gene expression data derived from patients diagnosed with various types of cancer. Our research objective revolves around acquiring insights into the optimal combination of dosages for chemotherapy  $(T)$ . In the domain of cancer treatment, the concurrent administration of different drugs has emerged as a prominent clinical approach [82]. In this context, patient survival serves as the designated outcome variable  $(Y)$ .
We utilize the dataset as employed in previous works [5, 67]. Specifically, we use the version from [5], which includes gene expression measurements of the 4,000 genes with the highest variability, serving as the features  $(X)$ , obtained from a cohort of  $n = 9659$  patients. Diverging from the approach of the aforementioned works, we once again model  $T$  as a combination of different dosages that are simultaneously assigned, taking into account their joint effect. This modeling choice closely aligns with medical practice, where the optimization of drug combinations for cancer therapy necessitates a holistic consideration of each patient's unique profile [16, 82].
We summarize the characteristics of the data in Table 3.
Variable	MIMIC	TCGA
Dim.	Meaning	Dim.	Meaning
X	33	patient covariates (age, sex, respiratory and cardiac measurements)	4000	gene expression measurements
T	2	respiratory rate, tidal volume	2	chemotherapy drugs (e.g., doxorubicin, cyclophosphamide)
Y	1	chance of patient survival	1	chance of patient survival
Table 3: Data characteristics of our standard settings.
Semi-synthetic data simulation: Our data simulation process is inspired by [5]. We use the observed features  $x$  from the real-world datasets. Then, for each treatment dosage  $j = 1, \ldots, p$ , we randomly draw two parameter vectors  $v_{1}^{(j)}, v_{2}^{(j)}$  from a normal distribution  $N(0,1)$  and normalize them via  $v_{1}^{(j)} = v_{1}^{(j)} / ||v_{1}^{(j)}||$  and  $v_{2}^{(j)} = v_{2}^{(j)} / ||v_{2}^{(j)}||$ , respectively. We define  $\eta^{(j)} = \frac{v_{1}^{(j)^T}x}{2v_{2}^{(j)^T}x}$  and
Available at: https://mimic.mit.edu
Available at: https://www.cancer.gov/ctga
Available at: https://github.com/ioanabica/SCIGAN
$\tilde{t}^{(j)} = \left(20 + 20\exp (-\eta^{(j)})\right)^{-1} + 0.2.$  We then sample the  $j$  -th dosage of the assigned dosage combination  $t$  via
$$
t ^ {(j)} \sim \operatorname {B e t a} (\tilde {\alpha}, q ^ {(j)}), \quad \text {w h e r e} \tilde {\alpha} = \alpha + 1, \quad q ^ {(j)} = \frac {\tilde {\alpha} - 1}{\tilde {t} ^ {(j)}} - \tilde {\alpha} + 2, \tag {12}
$$
with dosage bias  $\alpha = 2$ , if not specified otherwise. Hence,  $t^{(j)} \in [0,1]$  can be interpreted as a percentage or normalized version of the modeled dosages.
We define the dose-response function as
$$
\begin{array}{l} \mu (t, x) = 2 + \frac {2}{p} \sum_ {j = 1} ^ {p} \left(\left(1 + \exp (- \eta^ {(j)})\right) ^ {- 1} + 0. 5\right) \cdot \cos \left(3 \pi (t ^ {(j)} - \tilde {t} ^ {(j)})\right) - 0. 0 1 \left(t ^ {(j)} - \tilde {t} ^ {(j)}\right) ^ {2} \\ - 0. 1 \kappa \prod_ {j = 1} ^ {p} \left(t ^ {(j)} - \tilde {t} ^ {(j)}\right) ^ {2}, \tag {13} \\ \end{array}
$$
where  $\kappa = 1$  is an interaction parameter indicating the strength of the joint effect of the different dosages. Finally, we sample the observed outcome from  $y\sim N(\mu (t,x),0.5)$ . In our data simulation, we ensure that  $\tilde{t} = \left(\tilde{t}^{(j)}\right)_{j = 1}^{p}$  is the optimal dosage combination and that  $\tilde{t}$  additionally is the mode of the joint conditional dosage distribution of  $t\mid x$  with the generalized propensity score as density given by  $f(t,x) = \prod_{j = 1}^{p}\mathrm{Beta}(\tilde{\alpha},q^{(j)})$ .
Taken together, we adapt to a realistic setting in medicine where the dosage assignment is already informed to a certain degree but still can be further optimized. With increasing  $\alpha$ , the probability mass around the optimal dosage combinations increases, while, in areas further away from the optimum, overlap becomes more limited, and, hence, the estimation uncertainty of the dose-response function increases. We can interpolate between a random uniform distribution for  $\alpha = 0$  and a completely deterministic dosage assignment for  $\alpha = \infty$ .
Robustness checks: We now repeat the robustness checks from Sec. 5 also for the MLP baselines. The results are shown in Fig. 4 and Fig. 5. We observe that our method (DCNet+reliable) performs constantly well with respect to low regret and low (to no) variability. This holds true across the different settings. Also, we see that the baseline MLP clearly benefits from our reliable optimization during policy learning. Hence, the constrained optimization that we formulated in Steps 2 and 3 of our method are also a source of gain for existing baselines.
(b) TCGA
Figure 4: Robustness for dosage bias  $\alpha$ . Shown: selected policy (line) and the range over 5 runs (area). The naïve baseline has a large variability across runs while ours is highly robust as expected (i.e., we see no variability).
(a) MIMIC-IV
Figure 5: Robustness for number of dosages  $p$ . Shown: selected policy (line) and the range over 5 runs (area). The naïve baseline has a large variability across runs while ours is highly robust as expected (i.e., we see no variability).
High-dimensional dosage setting: We perform additional experiments to examine how our method scales to treatments of much higher dimensions. To improve the scalability of our method, we apply multiple prediction heads, each taking 3 treatment dimensions as input. Thereby, we fix the maximum number of jointly interacting dimensions to 3. Our motivation is that applying our method without adjustments to high treatment dimensions would lead to a high computational complexity of our prediction head because we would need to train weights for all possible interactions in the tensor product bases. However, it is often reasonable for real-world medical settings that the effective number of dimensions with joint effects is usually lower-dimensional, which is thus captured by our tailored approach. We report the results in Fig. 6. We observe that, unlike the other baselines, our method (DCNet+reliable) stays robust with low regret and low variation even for high dimensions.
Performance across multiple re-runs: To better assess the robustness of our method, we re-run our experiments for multiple train/val/test splits. Here, we also retrain the nuisance functions  $\hat{\mu}$  and  $\hat{f}$  for each split. We display the aggregated results for (i) data generated by our standard semi-synthetic
(a) MIMIC-IV
Figure 6: Robustness for high-dimensional number of dosages  $p$ . Shown: selected policy (line) and the range over 5 runs (area). Lower regret = better.
data simulation (Appendix B) in Table 4, and (ii) a more complex setting with a multimodal GPS in Table 5. For (ii), we replace the dosage combination assignment from Eq. (12) with a multimodal mixture density, i.e.,
$$
t ^ {(j)} \sim w ^ {(j)} \cdot \operatorname {B e t a} \left(\tilde {\alpha}, q ^ {(j)}\right) + \left(1 - w ^ {(j)}\right) \cdot \operatorname {B e t a} \left(\tilde {\alpha}, u ^ {(j)}\right), \tag {14}
$$
$$
w ^ {(j)} \sim \operatorname {B e r n o u l l i} (0. 5), \quad \tilde {\alpha} = \alpha + 1, \quad q ^ {(j)} = \frac {\tilde {\alpha} - 1}{\tilde {t} ^ {(j)}} - \tilde {\alpha} + 2, \quad u ^ {(j)} = \frac {\tilde {\alpha} - 1}{1 - \tilde {t} ^ {(j)}} - \tilde {\alpha} + 2.
$$
Methods	MIMIC-IV	TCGA
μ	π	selected	mean	std	selected	mean	std
Oracle (μ)	observed	1.089 ± 0.067	-	-	0.952 ± 0.061	-	-
MLP	naive	1.553 ± 0.990	1.391 ± 0.663	0.797 ± 0.430	0.832 ± 1.067	1.063± 1.176	0.303 ± 0.224
MLP	reliable	0.748 ± 1.516	0.849 ± 1.539	0.223 ± 0.375	1.080 ± 1.542	1.089 ± 1.542	0.022 ± 0.028
VCNet	naive	2.715 ± 0.938	2.322 ± 0.556	0.890 ± 0.337	3.169 ± 0.588	2.153 ± 0.406	0.935 ± 0.242
VCNet	reliable	0.033 ± 0.027	0.110 ± 0.183	0.171 ± 0.377	0.044 ± 0.017	0.114 ± 0.145	0.159 ± 0.354
DCNET	naive	2.512 ± 0.787	0.974± 0.700	1.134± 0.336	2.048 ± 0.775	1.682 ± 0.478	0.776 ± 0.263
DCNET	reliable (ours)	0.025 ± 0.021	0.024 ± 0.020	0.001 ± 0.001	0.025 ± 0.010	0.263 ± 0.540	0.337 ± 0.750
Reported: selected / mean / standard deviation (lower = better)
Table 4: Regret of 5 re-runs (mean ± std) with different train/val/test splits with each  $k = 5$  restarts for policy learning.
Methods	MIMIC-IV	TCGA
μ	π	selected	mean	std	selected	mean	std
Oracle (μ)	observed	1.294 ± 0.156	-	-	1.624 ± 0.071	-	-
MLP	naive	2.139 ± 1.539	0.860 ± 0.566	0.993 ± 0.447	1.745 ± 1.270	1.444 ± 1.432	0.640 ± 0.166
MLP	reliable	1.069 ± 1.518	0.544 ± 0.509	0.786 ± 0.656	0.465 ± 0.902	0.720 ± 1.488	0.157 ± 0.328
VCNet	naive	1.885 ± 0.258	1.627 ± 0.149	0.443 ± 0.152	1.728 ± 0.306	1.659 ± 0.176	0.368 ± 0.130
VCNet	reliable	0.126 ± 0.065	0.192 ± 0.140	0.147 ± 0.193	0.029 ± 0.027	0.119 ± 0.117	0.201 ± 0.279
DCNET	naive	0.713 ± 0.801	0.389 ± 0.256	0.515 ± 0.347	0.260 ± 0.409	0.141 ± 0.125	0.209 ± 0.194
DCNET	reliable (ours)	0.077 ± 0.059	0.065 ± 0.038	0.061 ± 0.058	0.014 ± 0.010	0.028 ± 0.026	0.031 ± 0.060
Reported: selected / mean / standard deviation (lower = better)
Table 5: Regret of 5 re-runs (mean ± std) with different train/val/test splits with each  $k = 5$  restarts for policy learning in an additional setting simulated with a multimodal GPS.
We can interpret the results as follows:
(1) The column "selected" displays the final performance of our method when using the finally recommended policy. In this column, the mean  $\pm$  std can be considered as the expected regret with Monte-Carlo-CV confidence intervals (including retraining the nuisance functions) and should primarily be used for evaluating the performance and for comparison with different policy learning methods. We observe that our method significantly outperforms the baselines.
(2) The column "mean" represents the means within the  $k$  restarts of the policy learning (i.e., Step 3 in our method). As such, its values can indicate how a randomly selected policy out of the  $k$  restarts would perform compared to the "selected one". For example, we observe that, for the "naive" baselines using a naive MSE loss for selection, suboptimal policies are oftentimes selected, whereas our method shows robust behavior.
(3) The column "std" shows the standard deviation within the  $k$  restarts of the policy learning (i.e., Step 3 in our method). As such, the values indicate how often similar performing policies are learned. It thus refers to the robustness within the  $k$  runs. As a consequence, lower values imply that (i) probably fewer restarts  $k$  are needed to find the targeted optimum, and (ii) fewer local optima are learned, which reduces the risk of selecting a suboptimal policy.
Hence, we suggest using the "selected" column (mean ± std) for evaluating the final performance of our method, and the other two columns ("mean" and "std") to show the inner robustness of our method compared to ablation baselines.
Overall, we demonstrate that the performance of our framework is robust across the tested settings and has low regret and low variation across different re-reruns and different restarts, including the more complex setup.
Sensitivity analysis of the reliability threshold  $\bar{\varepsilon}$ : In our experiments, we set the reliability threshold  $\bar{\varepsilon}$  to the  $5\%$  -quantile of the estimated GPS  $\hat{f} (t,x)$  of the train set. This is due to the reason that the GPS is a probability density function over a  $p$  -dimensional continuous space and, hence, is hard to interpret intuitively. Unlike discrete probabilities, which are bounded between 0 and 1 and can be interpreted intuitively by humans, the GPS has no upper bound. To this end, we suggest a heuristic to choose  $\bar{\varepsilon}$  dependent on  $x\%$  -quantiles of the estimated GPS based on the train dataset. As a result, we aim to learn optimal dosages, which have at least the same estimated reliability as  $x\%$  of the observed data.
In our heuristic, the selection of the quantile  $x$  is still subjective and cannot be optimized during hyperparameter tuning, as we assume to have no access to the ground truth  $\mu(t,x)$ . Nevertheless, we can leverage our semi-synthetic data to perform a sensitivity analysis over  $x$  for  $\bar{\varepsilon}$  and thus corroborate our choice. For this, we evaluate the final performance of our method on the test data for  $\bar{\varepsilon}$  set to different quantiles of  $\hat{f}$  over  $k = 5$  runs. The results are displayed in Fig. 7.
We observe that, depending on the dataset, different quantiles for selecting  $\bar{\varepsilon}$  can lead to differences in variability. However, we find that lower values tend to lead to less variability. Also, we find that, even when there is higher variability within the  $k$  runs, for every setting, a run with low regret is selected. We attribute this to our validation loss from Eq. 9. Together, the results are in line with our expectations and add to the robustness of our methods.
Figure 7: Sensitivity analysis of the reliability threshold  $\bar{\varepsilon}$  for different quantiles of  $\hat{f}$  of the train set. Shown: selected policy (line) and the range over 5 runs (area).
D Insights
Individualized dose-response estimates: In our experiments in Sec. 5 and in Appendix C, we showed a comparison to naïve baselines. Therein, in particular, we mainly demonstrated that the reliable optimization step in our method is an important source of performance gain. We now give additional intuition as to why also our DCNet is a source of performance gain. In the following, we first demonstrate that our DCNet is able to estimate the dose-response function more accurately than the baseline in regions with sufficient overlap. We then discuss why this brings large gains for our reliable optimization. Thereby, we offer an explanation why both DCNet and reliable policy optimization are highly effective when used in combination.
In Table 6, we display the performance of the dose-response estimators. In detail, we report: (i) the mean squared error ("Biased MSE"), calculated on a holdout test dataset which follows the same biased dosage assignment as the train data; (ii) the mean ("MISE", see e.g., [67]) of the integrated squared error calculated on areas with sufficient overlap, i.e.  $\sum_{i=1}^{n} \int_{[0,1]^p} (\mu(t, x_i) - \hat{\mu}(t, x_i))^2 dt \mid \hat{f}(t, x) > \bar{\varepsilon}$ ; and (iii) the standard deviation of the integrated squared error on areas with sufficient overlap ("SD-ISE"). Intuitively, the MISE and SD-ISE are the expected value and variation of the prediction error of the dose-response estimator over the covariate-treatment space with sufficient overlap. Hence, a low value implies that all suitable dosage combinations with sufficient overlap are estimated properly and that certain, less frequent or non-dominant combinations are not neglected. We observe that our DCNet also performs clearly best wrt. to all metrics, undermining its ability to successfully account for the different drug-drug interactions under dosage bias.
Methods	MIMIC-IV	TCGA
μ	Biased MSE	MISE	SD-ISE	Biased MSE	MISE	SD-ISE
MLP	0.293	0.049	0.069	0.438	0.374	0.745
VCNet	0.691	0.558	0.749	0.551	0.406	0.474
DCNet (ours)	0.260	0.013	0.032	0.257	0.006	0.010
Table 6: Estimation error of dose-response estimators.
To give further intuition to the above, we select one random observation per dataset and plot the following functions: (i) the estimated GPS, (ii) the oracle individualized dose-response surface, (iii) the estimates using DCNet, and (iv) the estimates using the baseline MLP. We compare two settings, namely with no dosage bias  $(\alpha = 0)$  and thus sufficient overlap (in Fig. 8) and with dosage bias  $(\alpha = 2)$  and thus clearly limited overlap (in Fig. 9).
The plots can be interpreted as follows. With no dosage bias (Fig. 8), both DCNet and MLP seem to approximate  $\mu$  adequately (with slightly better estimates of DCNet at the modes at the boundary areas) and are thus suitable for naïve policy learning. However, the findings change in a setting with limited overlap (Fig. 9). Here, the limited overlap is indicated by areas with low estimated GPS  $\hat{f}$ . Both estimators for  $\hat{\mu}$  show clear deviations to the oracle  $\mu$ . However, the deviations are smaller for large parts in the case of our DCNet (as compared to the MLP). Evidently, the dose-response surface predicted by the MLP loses its expressiveness, especially for the high-dimensional TCGA dataset, and fails to model the maximum adequately. This leads to erroneous predictions when used as a plug-in estimator for policy learning, which can not be compensated by constraining the search space to regions with sufficient overlap. DCNet, on the other hand, is able to model the dose-response surface more adequately in areas with sufficient overlap due to its custom prediction head ensuring expressiveness. However, in regions with limited overlap, the splines of the prediction head extrapolate and lead to strongly incorrect predictions. Anyway, as our goal is to learn a reliable policy, we aim to exclude the regions with limited overlap and thus high uncertainty (yellow regions) such that they do not affect our final predicted policy.
In conclusion, when now applying our reliable policy optimization on top of our DCNet, we constrain the search space to the reliable regions, in which DCNet outperforms the MLP. Thus, we ensure reliable and accurate predictions for policy learning by using our method.
(b) TCGA
Figure 8: Insights for one randomly sampled observation in the setting with dosage bias  $\alpha = 0$ . We show (i) the estimated GPS and, additionally, the individualized dose-response surface through (ii) the true oracle function, (iii) the estimates using our DCNet, and (iv) the estimates using the baseline MLP. The optimal dosage combination is shown as a red point.
(b) TCGA
Figure 9: Insights for one randomly sampled observation in the setting with dosage bias  $\alpha = 2$ . We show (i) the estimated GPS and, additionally, the individualized dose-response surface through (ii) the true oracle function, (iii) the estimates using our DCNet, and (iv) the estimates using the baseline MLP. The optimal dosage combination is shown as a red point.
Computational complexity of policy optimization: To ensure scalability at inference time, we train an additional policy neural network to directly predict the optimal policy, instead of performing optimization for new patients. We explain our motivation in the following: The complexity of a direct optimization at inference time is non-trivial and depends on the complexity of the forward pass of DCNet  $\sigma_{\mu}$  and the CNFs  $\sigma_f$ , as well as the selected constrained optimization algorithm, and the number of patients to evaluate  $n_{\mathrm{eval}}$ . While the exact complexity of the direct optimization depends on the specific solver cost  $\sigma_{\mathrm{solv}}$ , which, for most algorithms, scales heavily with the dosage dimension  $p$ , it is at least  $\mathcal{O}(n_{\mathrm{eval}} \times (\sigma_{\mu} + \sigma_f) \times \sigma_{\mathrm{solv}})$ . This becomes costly for large-scale datasets in practice. In contrast, our policy network is trained only once. Hence, upon deployment, it simply needs operations for evaluation, which is at most  $\mathcal{O}(n_{\mathrm{eval}} \times \sigma_{\mathrm{net}})$ , where  $\sigma_{\mathrm{net}}$  is the complexity of only the forward pass of the policy network. Hence, it usually holds  $\sigma_{\mathrm{net}} \ll \sigma_{\mathrm{solv}}$ .
E Modeling and Training Details