Deep Reinforcement Learning for Cost-Effective Medical Diagnosis
Authors: Zheng Yu
Yikuan Li
Joseph Chahn Kim
Kaixuan Huang
Yuan Luo
Mengdi Wang
Venue: ICLR 2023 poster (2023)
Keywords: medical diagnostics
Pareto front
reinforcement learning
non-Markovian reward
semi-model-based policy optimization

Abstract: Dynamic diagnosis is desirable when medical tests are costly or time-consuming. In this work, we use reinforcement learning (RL) to find a dynamic policy that selects lab test panels sequentially based on previous observations, ensuring accurate testing at a low cost. Clinical diagnostic data are often highly imbalanced; therefore, we aim to maximize the F1 score instead of the error rate. However, optimizing the non-concave $F_1$ score is not a classic RL problem, thus invalidating standard RL methods. To remedy this issue, we develop a reward shaping approach, leveraging properties of the $F_1$ score and duality of policy optimization, to provably find the set of all Pareto-optimal policies for budget-constrained $F_1$ score maximization. To handle the combinatorially complex state space, we propose a Semi-Model-based Deep Diagnosis Policy Optimization (SM-DDPO) framework that is compatible with end-to-end training and online learning. SM-DDPO is tested on diverse clinical tasks: ferritin abnormality detection, sepsis mortality prediction, and acute kidney injury diagnosis. Experiments with real-world data validate that SM-DDPO trains efficiently and identify all Pareto-front solutions. Across all tasks, SM-DDPO is able to achieve state-of-the-art diagnosis accuracy (in some cases higher than conventional methods) with up to $85\%$ reduction in testing cost. Core codes are available at https://github.com/Zheng321/Deep-Reinforcement-Learning-for-Cost-Effective-Medical-Diagnosis.

Published as a conference paper at ICLR 2023
DEEP REINFORCEMENT LEARNING FOR COST-EFFECTIVE MEDICAL DIAGNOSIS
Zheng Yu" Princeton University
Yikuan Li
Joseph C. Kim Princeton University
Kaixuan Huang" Princeton University
Northwestern University
Yuan Luo
Mengdi Wang!
Northwestern University
Princeton University
ABSTRACT
Dynamic diagnosis is desirable when medical tests are costly or time-consuming. In this work, we use reinforcement learning (RL) find a dynamic policy that selects lab test panels sequentially based on previous observations, ensuring accurate testing at low cost. Clinical diagnostic data are often highly imbalanced: therefore, we aim to maximize the Fl score instead of the eπrоr rate. However, optimizing the non-concave Fi score not a classic RI. problem, thus invalidating standard RI. methods. To remedy this issue, we develop reward shaping approach, leveraging properties of the F1 score and duality of policy optimization, to provably find the set of all Pareto-optimal policies for budget-constrainee Fi score maximization. To handle the combinatorially complex state space, we propose a Semi-Model-basee Deep Diagnosis Policy Optimization (SM-DDPO) framework that is compatible with end-to-end training and online learning. SM-DDPO is tested on diverse clinical tasks: ferritin abnormality detection, sepsis mortality prediction, and acute kidney injury diagnosis. Experiments with real-world data validate that SM-DDPO trains efficiently and identify all Pareto-front solutions. Across all tasks, S.M-DDPO is able to achieve state-of-the-ar diagnosis accuracy (in some cases higher than conventional methods) with up to 85% reduction in testing cost. Core codes are available on GitHub!
INTRODUCTION
In clinical practice, physicians usually order multiple panels of lab tests on patients and their inter- pretations depend medical knowledge and clinical experience. Each test panel is associated with certain financial cost. For lab tests within the same panel, automated instruments will simultaneously provide all tests, and eliminating single lab test without eliminating the entire panel may only lead to small reduction in laboratory cost (Huck& & Lewandrowski 2014). On the other hand, concurrent lab tests have been shown exhibit significant correlation with each other, which can be utilized to estimate unmeasured test results (Luo et al.) 2016). Thus, utilizing the information redundancy among lab tests can be promising way of optimizing which test panel order when balancing comprehensiveness and cost-effectivenes. The efficacy of the lab test panel optimization can be evaluated by assessing the predictive power of optimized test panels on supporting diagnosis and predicting patient outcomes.
We investigate the use of reinforcement learning (RL) for lab test panel optimization. Our goal is to dynamically prescribe test panels based on available observations, in order to maximize diagno- sis/prediction accuracy while keeping testing at low cost. is quite natural that sequential test panel selection for prediction/Classifcatiin can be modeled as a Markov decision process (MDP).
CCo-first authors.
CCosennio authors.
Contact information: 7. Yu, Kim, K. Huang and M. (zhengy, josephck, ka xuanh, mengdiw w wwwpinnctoonn.oom edu; Y. Li and Y. Luo: (yi kuan. 1i, yuan. luo)

---

Published as a conference paper at ICLR 2023
However, application of reinforcement learning (RL) to this problem is nontrivial for practical considerations. One practical challenge is that clinical diagnostic data are often highly imbalanced, in some cases with <5% positive cases (Khushi et al. 2021 Li et al. 2010: Rahman & Davis 2013). In supervised learning, this problem is typically addressed by optimizing towards accuracy metrics suitable for unbalanced data. The most prominent metric used by clinicians is the F1 score, i.e., the harmonic mean of a prediction model's 's recall and precision, which balances type I and type II errors in single metric. However, the Fl score not a simple weighted error rate- this makes designing the reward function hard for RL. Another challenge is that, for cost-sensitive diagnostics, one hopes to view this as a multi-objective optimization problem and fully characterize the cost-aceuracy tradeoff, rather than finding an ad-hoc solution on the tradeoff curve. In this work, we aim to provide tractable algorithmic framework, which provably identifies the set of all Pareto-fron policies and trains efficiently. Our main contributions are summarized as follows:
We formulate cost- sensitvee diagnosties as a multi-objective policy optimization problem. The goal is to find all optimal policies the Pareto front of the cost-accuracy tradeoff.
To handle severely imbalanced clinical data, we focus on maximizing the Fi score directly. Note that F1 score is nonlinear, nonconvex function of true positive and true negative rates. It cannot be formulated as a simple sum of cumulaative rewards, thus invalidating standard RL solutions. We leverage monotonicity and hidden minimax duality of the optimization problem, showing that the Pareto set can be achieved via a reward shaping approach.
We propose Gemi-Model-based Deep Diagnostic Policy Optimization (SM-DDPO) method for learning the Pareto solution set from clinical data. Its architecture comprises three modules and can be trained efficiently by combing pretraining, policy update, and model-based RL.
We apply our approach to real-world clinical datasets. Experiments show that our approach exhibits good accuracy-cos trade-off on all tasks compared with baselines. Across the experiments, our method achieves state-of-the-an accuracy with up to 80% reduction in cost. Further, SM-DDPO able to compute the set of optimal policies corresponding to the entire Pareto front. We also demonstrate that SM-DDPO applies not only to the F1 score but also to alternatives such as the AM score.
2 RELATED WORK
Reinforcement learning (RL) has been applicd in multiple clinical care settings to learn optimal treatment strategies for sepsis Komorowski et al. (2018). to customize antiepilepsy drugs for seizure control Guez et al. 2008) etc. See survey Yu et al. (2021) for more comprehensive summary. Guidelines on using RL for optimizing treatments in healthcare has also been proposed around the topics of variable availability, sample size for policy evaluation, and how to ensure learned policy works prospectively as intended Gottesman et al. (2019). However, using RL for simultaneously reducing the healthcare cost and improving patient's 's outcomes has been underexplored.
Our problem of cost-sensitive dynamic diagnosis/prediction is closely related to feature selection in supervised learning. The original static feature selection methods, where there exists common subset of features selected for all inputs, were extensively discussed Guyon Elisseeff (2003); Kohavi &
John (1997); Bi et al. (2003); Weston et al. (2003) 2000). Dynamic feature selection methods He al. 2012); Contardo al.] (2016): Karayev et al. (2013), were then proposed to take the difference between inputs into account. Different subsets of features are selected with respect to different inputs.
By defining certain information value of the features Faay Yang (2019): Bilgic Getoor (2007). or estimating the gain of acquiring a new feature would yield Chai et al. (2004). Reinforcement learning based approaches Ji & Carin (2007); Trapeznikov & Saligramal (2013); Janisch et al. (2019): Yin et al.]
(2020): Li & Oliva (2021) y Nam et al. (2021) are also proposed to dynamically select features for prediicion/ccasiffcatioo We give a more detailed discussion in Appendix A
3 PARETO-FRONT PROBLEM FORMULATION
3.1 MARKOV DECISION PROCESS (MDP) MODEL
We model the dynamic diagnosis/prediction process for new patient as episodic Markov decision process (MDP) M = (S,A, R, E). As illustrated in Figure I the state of a patient is described by M, where Rd denotes d medical tests of a patient, M {0, 1)d is binary mask
2

---

Published as a conference paper at ICLR 2023
indicating whether the entries of observed missing. Let there be D test panels, whose union the set of all d tests. The action set A = ( .D Li [P,N] contains two sets of actions observation actions and prediction/diagnosis actions. At each stage, one either pick an action E {1, 2, ........................... D} from any one of the available panels, indicating choosing test panel to observe, which will incur corresponding observation cost c(a); Or can terminate the episode by directly picking a prediction action a E [P. N}, indicating diagnosing the patient as positive class (P) or negative class (N). A penalty will generated if the diagnosis does not match the ground truth y. An example of this process in sepsis mortality prediction is illustrated in Figure [] We considers the initial distribution & be patients with only demographics panel observed and discount factor r
Falle
CM?
pandls at limistap
Inau
obsoration addent
dbwieed yer nanele Stredp:
pullon- novol noilly
Erd epticce
trnongg 2
nair diggosse revard corrett &igg
CMP
ABG
AFTT undieredd pEMES
APTT
APIT
undisseredd pares Emrerpp
undioreel pando: STedEs
limette
Figure 1: MDP model of dynamic diagnosis: illustration of state state-action transitions in one episode.
3.2 MULTI-ORJECTIVE POLICY OPTIMIZATION FORMULATION
Let S - AA be the overall policy, map from the set of states the set of actions. We optimize towards two objectives:
Maximizing prediction accuracy. Due to severely imbalanced data, we choose to maximize the Fi score denoted by Fi (T), as a function of policy æ. Fi score measures the performance of the diagnosis by considering both type and type II errors, which is defined as:
TP(x) TP(T)
2TP(m)
(x
(FP(F) + FN(T)) 1+ TP(x) TN((T)
where TP(T), (P(T FN(TT) are normalized true positive, true negative, false positive and false negative that sum up to Remark that TP(T), TN(x), FP(z), (P((),,N() FN(x)) can all expressed as sum of rewards/costs over the MDP's state trajectories. However, (Ir)) is nonlinear with respect to the MDP's 's state-action occupancy measure, thus cannot be expressed as any cumulative sum of rewards.
Lowering cost. Define the testing cost by Bost(T)=BB*I ExeLDI c(k) 1fa, =k]], where =
E" denotes expectation under policy π, c(k) is the cost of panel In this work, we hope to solve for cost-sensitive policies for all possible testing budget. In other words, we aim to find the cost-sensitive Pareto front as follows.
Definition 3.1 (Cost-Fi Pareto Front of Multi-Ohjective Policy Optimization). The Pareto front IT" for sest-iensi dynamic diagnosis/predictinn is the set of policies such that
r UBs0 argmax { Fi (T)) subject Cost( S <B}
Finding IT" requires novel solutions beyond standard RL methods. Challenges are two-folded: (1) Even the single-objective case, (T) is a nonlinear, non-concave function of TP(T), TN((). Although both TP(#), TN(=) can be formulated as expected sum of rewards in the MDP, the Fi score is never simple sum of rewards. Standard RL methods do not apply to maximizing such function. (2) We care about finding the set of all Pareto-optimal policies when there are two conflicting objectives, rather than an ad hoc point on the trade-off curve.
2 An alternative to the Fl score is the AM metric that measures the average of true positive rate and true negative rate for imbalanced data Natarajan al. (2018) Menon et al. (203)). Our approach directly applies to such linear metrie. Please refer Appendix for details.
3

---

Published as a conference paper at ICLR 2023
4 FINDING COST-F, PARETO FRONT VIA REWARD SHAPING
The Fi score is a nonlinear and nonconvex function of true positive, true negative, false positive, false negative rates. It cannot be expressed by sum of rewards. This invalidates all existing RL methods even in the unconstrained case, creating tremendous challenges.
Despite the non-concavity and nonlinearity of F1, we will leverage the mathematical nature of Markov decision process and properties of the Fl score to solve problem (1). In this section, we provide an optimization duality analysis and show how to find solutions problem (1) via reward shaping and solving reshaped cumulative-rewart MDP.
Step 1: utilizing monotonicity of F1 score To start with, we note that Fi score monotonically increasing in both TP and TN. Assume, for any given cost budget B, the optimal policy π² (B) achieves the highest F1 score. Then r= (B) is also optimal to the following program:
max (T) subject Cost ( s B, TP(x) (B()
indicating the Pareto front of F1 score a subset of
II" C [nso.kelo..] argmax (TN(T) subject to Cost(x) B, (P(TT 2 K}.
(2)
Step 2: reformulation using occupancy measures Fix any specific pair B*). Consider the equivalent dual linear program form Zhang et al. (2020) of the above policy optimization problem 2
It in terms of the cumulative state- action occupancy measure μ A + S defined as:
20
1² (s. a) = ET E1 1s: = s, s,a, = a)} Vs E S,a E A. Then the program is equivalent to:
max TN( subject to Cost(p)= S B, TP()) 2 K, Eu(s,a)
E
u(s', P(s a')) + E(s), Vs
s',a'[[D]
where 5(-) denotes initial distribution, and TP, TN and cost are reloaded in terms of occupancy as:
TP(p)= = Σ a)s,a), TN(R) = Σ ((s,a), Cost(p)= = Σ c(k) Σ M(s,a).
p=P,==P
y=N,a=N
kE(D)
Step 3: utilizing hidden minimax duality The above program can be equivalently reformulated as a max-min program:
max miH
TN(#) + + +A. (TP(M) K) p. (Cost(n) B)
A20.P50
subject Σ)(,,o)= Σ m(s', d') a') E(s),
Note the max-min objective is linear in terms of X, and μ.. Thus, minimax duality holds, then we can swap the min and max to obtain the equivalent form:
min max TN(H) + (TP(p) K) p. (Cost(n) B) x20,e00
subject to Em(,,o) Σ m(s', a''',a')Pi P'ssl'' a') Ess),
For any fixed pair of (A, p). the inner maximization problem of the above can be rewritten equivalently into an unconstrained policy optimization problem: max, TN(x) + X (P(TT) + p Cos This is finally a standard cumulative-sum MDP problem, with reshaped reward: reward p c(t) for the action of choosing test panel I, reward À for the diagnosis action and get true positive, reward for getting a true negative. Putting together three steps, we can show the following theorem. The full proof can be found in Appendix
Theorem 4.1. The Cost-Fi Fi Pareto front defined in 1 a subset of the collection of all reward-shaped
solutions, given by
IT' C II Ux20, argmax (TN(T) + -P((P)() + p-Cos((T)
Thus, to learn the full Pareto front, it suffices to solve a collection of unconstrained policy optimization problems with reshaped cumulative rewards.
4

---

Published as a conference paper at ICLR 2023
5 METHOD
In this section, we propose a deep reinforcement learning pipeline for Pareto-optimal dynamic diagnosis policies. We use modular architecture for efficient encoding of partially-observe patient information, policy optimization and reward learning.
5.1 ARCHITECTURE
Our Momi-Model-based Deep Diagnostic Policy Optimization (SM-DDPO) framework is illustrated in Figure 2 The complete dynamic testing policy comprises three models: (1) posterior state encoder for mapping -obrerll patient information to an embedding vector; (2) state-to- diagnosis/prediction classifier which can be reviewed as a reward function approximator; (3) test panel selector that outputs action based on the encoded state. This modular architecture makes RL tractable via a combination pre-training, policy update and model-based RL.
State
Demegraphics Med test
M
Semb
Panel/prediction selector
Action: select a test panel diagnose
Reward Estimate
Posterior State Encoder
Classifler fe
Diagnose
CrossEntropy
Update
end-to-and
training
Figure 2: Dynamic diagnostic policy learning via semi-modee bbaeed proximal policy optimization. The full policy и comprises of three modules: posterior state encoder, classifier, and panel selector.
5.2 POSTERIOR STATE ENCODER
We borrowed the idea of imputation map the partially observed patient information posterior embedding vector. In this work, we consider flow-based deep imputer named EMFlow Given the imputer (-) parameterized by 0, the RL agent observes tests X and calculates Imp, (xOMM)
Rd posterior state encoder. Unlike conventional imputation Lin & Tsai (2020); Austin al. (202T) Osman et (2018).. our posterior state encoder aims at resolving exponentially many possible missing patterns. Therefore, we pretrain it unlabeled augmented data, constructed by repeatedly and randomly masking entries as additional samples.
5.3 END-TO-END TRAINING VIA SEMI-MODEL- BASED POLICY UPDATE
Training the overall policy using a standard RL algorithm alone (such as Q learning, or policy gradient) would suffer from the complex state and action spaces. To ease the heavy training, design semi-model-basse modular approach to train the panel selector and classifier concurrently but different manners:
The classifier fa(-) Rd + I2, parameterized by é. maps the posterior encoded state Impa O M) probability distribution over labels. It trained by directly minimizing the cross entropy loss f. from collected dat= This differs from typical classification in that the data are collected adaptively by RL, rather than sampled from a prefixed source.
The panel selector, network module parameterized by v. takes the following as input
emb emh x M) (Impo(x @ (Imp M)),M) S
The EMFlow imputaaion method is originally proposed in Ma Ghosh (022) that maps the data space to
a Gaussian latent space via normalizing flows. We give more detailed discussion of this method in Appendix "The forms of the training objective f. and of classifier and panel selector given in Appendix D
5

---

Published as a conference paper at ICLR 2023
and maps it to probability distribution actions. We train the panel selector by using the classical proximal policy updates (PPO) Schulman et (017), which aims at maximizing a clipped surrogate objective regularized by square-erroo loss of the value functions and entropy bonus. We denote this loss function as Ett and relegate its expanded form to Appendix D
The full algorithm updates panel selector and classifier concurrently, given in Algorithm and visualized in Figure 2 We call it aeeii-model-bee
because it maintains a running estimate of the classifier (which a part of the reward model of the MDP) while making proximal policy update.
Algorithm 1 Semi-Modell Deep Diagnosis Policy Optimization (SM-DDPO)
Initialize: Impg. Classifier falo.)) Pancl/Prediction Selection Policy Tu(o.0), number of loops L, L,, L2, stepsize r7. for i = 0, 1, ... I do
yemb End-to-end training outer loop
Construct RL environment using state embedding 0, defined (
for j = 1, ........ T, do
b Policy update inner loop
Run RI. policy in environment for T timesteps and save observations in
Update panel selection policy by vilh)) argmax (x);; (165111)
end for
Set (+++1,0)) 16(5.12)
for 1,2, .... Le do
b Classifier update inner loop
Sample minibatch B, from Q
Update classifier by Φ/. 26- -n7 -.VVVEE ExER, 100-11((0 Mx))
end for Set ((++1.0) g(1.12)
end for
Output: Classifier faur+1,0,, Policy Tg(1+100)
Such hybrid RL technique of model learning and policy updates has been used for solving complex games, where notable example is Deepmind's Muzero Schritwieesr et al.) (2020). Further, we remark that Algorithm does end-to-end training. Thus, compatible with on-the-fly learning. The algorithm can start with as little as zero knowledge about the prediction task, and it can keep improving on new incoming patients by querying test panels and finetuning the state encoder.
EXPERIMENTS
We test the method three clinical tasks using real-world datasets. See Table T for summary. We split each dataset into 3 parts: training set. validation set, and test set. Training data is further split into two disjoint sets, one for pretraining the state encoder and the other one for end-to-on RI training. Validation set is used for tuning hyperparameters During RI. training we sample a random patient and a random subset of test results as initial observations at the beginning of an episode, for sufficient exploration. We evaluate the trained RL policy patients from the test sets, initialized at a state with zero observed test result, and report Fl score and AUROC.
Table I: Summary statisties of ferritin, AKI and sepsis datasets
Dataset
# of tests test panels
patients
% positive class
#training
#validation
#held-out testing
Ferritin
43,472
8.9%
32,602
6,522
4,248
AKI
23,950
16.5% 14.5%
17.964
3,600
2,386
Sepsis
28
5,783
4,335
869
5579
6.1 CLINICAL TASKS
We briefly describe three clinical tasks for our experiments. We refer to Appendix B for more details.
Ferritin abnormality detection Blood ferritin level can indicate abnormal iron storage, which is commonly used to diagnose iron deficiency anemia Short & Domagalski (2013) or hemochromatosis
Detailed data splitting, hyperparameter choices and searching ranges are presented in Appendix B
'Our codes used the implementation of PPO algorithm in package Raffin al. 2021

---

Published as a conference paper at ICLR 2023
(iron overload) Crownover& & Covey (2013). Machine learning models can predict abnormal ferritin
levels using concurrent laboratory measurements routinely collected in primary Vare|Luo et al. (2016)
Kurstjens et al. (2022). ) These predictive models achieve promising results, e.g. around 0.90 AUC
using complete blood count and C-reactive protein Kurstjens et al. (2022). and around 0.91 AUC
using common lab tests Luo et al.] (2016). However, both studies required the full observation of all
selected predictors without taking the financial costs into consideration.
We applied our proposed models to a ferritin dataset from tertiary care hospital (approved by Institutional Review Board), following the steps described in Luo et al. (206)) Our dataset includes 43,472 patients, of whom 8.9% had ferritin levels below the reference range that should be considered abnormal. We expected to predict abnormal ferritin results using concurrent lab testing results and demographic information. These lab tests were ordered through and can he dynamically selected from 6 lab test panels, including basic metabolic panel (BMP, n=9, [estimated national average] cost=S36), comprehensive metabolic panel (CMP, n=16, cost=S48), basic blood count (BBC, n=10, cost=826), complete blood count (CBC, n=20, cost=$44), transferrin saturation (TSAT, n=2, cost=$40) and Vitamin B- test (n=1, cost=S66).
Acute Kidney Injury Prediction Acute kidney injury (AKI) is commonly encountered in adults in the intensive care unit (ICU), and patients with AKI are at risk for adverse clinical outcomes such as prolonged ICU stays and hospitalization, need for renal replacement therapy, and increased mortality Kellum & Lameire (201)). AKI usually occurs over the course of a few hours to days and the efficacy of intervention greatly relies the early identification of deterioration Kellum & Lameire (2013). Prior risk prediction models for AKI based on EHR data yielded modest performance, e.g.,
around 0.75 AUC using a limited set of biomarkers Perazella (2015) or a specific group of patients Sanchez-Pinto & Khemani (016)) or around 0.8 AUC using comprehensive lab panels of general IdultCCU populations such as from the MIMIC ZZimmermn al. (2019); (Sun et al. (2019).
In this experiment, we followed steps in Zimmerman et al. (2019) to extract 23,950 ICU visits of 19,811 patients from the MIMIC-III dataset Johnson et al. (2016), among which 16.5% patients develop AKI during their ICU stay. We aimed at predicting the AKI onset within 72 hours of ICU admission using a total of 31 features including demographics, physiologic measurements, and lab testing results extracted within 24 hours of ICU admission. The lab tests were categorized into 4 panels, i.e. CBC (n=3, cost=$44), CMP (n=8, cost=$48), the arterial blood gas panel (ABG, n=2, cost=$473) and the activated partial thromboplastin time panel (APTT, n=6, cost=$26). The demographic information and physiologic measurements were collected before test panel selection, thus they are considered visible.
Sepsis Mortality Prediction for ICU Patients Sepsis is a life-threatening organ dysfunction, and is leading cause of death and cost overruns in ICU patients Angus & Van der Poll (2013) Early identification of risk of mortality in septic patients is important to evaluate the patients' status and improve their clinical outcomes Moreno et al. (2008). Most of the previous sepsis mortality prediction models use all available test results as predictors, without balancing the cost of ordering all the
associated test panels Lee et al. p0 Ding & Luo (00T): Shin et al. (202T); Moreno et al. (2008).
We followed steps in Shin et al. 2022)) to collect 5,783 septic patients from the MIMIC-III dataset Johnson et al. (2016) according the Sepsis-3 criteria (Singer et al. (2016). The in-hospital mortality rate of this cohort is 14.5%. We focused on predicting in-hospital mortality for these sepsis patients using demographics information, medical histories, mechanical ventilation status, the first present lab testing results and physiologic measurements within 24 hours of ICU admission, and the Sequential Organ Failure Assessment (SOFA) score. Similarly to the setup in the AKI experiment, the lab tests were also categorized into 4 panels of CBC (n=5), CMBP (n=15), ABG (n=6) and APTT (n=2). The components of SOFA score may be based on the lab testing results in CBC. CMP or ABG panels Singer et al. (216)). Demographic features are considered visible.
6.2 PERFORMANCE RESULTS
Our method is tested with comparison to a number of baselines that use either full/partial test results or statically/dynamical selected tests for prediction. They include logistic regression, random forest Ho (1995). XGBoost Chen & Guestrin (2016), LightGBM Ke et al. (2117). 7) a 3-layer multi-layer perceptron, as well RL-based approach Janisch et (2019). Experiments results, such as Fl score, AUROC and testing costs are reported in Table We emphasize that these baselines are incapable to handle the task of finding the Pareto front. Thus, we only test them under no budget constraints.
7

---

Published as a conference paper at ICLR 2023
Table 2: Model performance, measured by Fj score, area under ROC (AUC), and testing cost, for three real-world clinical datasets. The tested models include logistic regression (LR), random forests (RF), gradient boosted regression trees CeBBosttChen & Guestrin (2016) and LightGBM Ke et al. (2017)
). 3-layer multi-layer perceptron, Q-learning for classification with costly features (CWCF) Janisch
ctt al. 2019), Random Selection (RS). Fixed Selection (FS). All models were fine-tuned to maximize the F1 score. The model yielded the highest F1 score is in bold. The model required the least testing
cost is underlined. More detailed results of this table with more dynamic baselines and standard
deviations reported in Appendix B
Models
Ferritin
AKI
Sepsis
Test Selection
Metries
Fi
AUC
Cost
Fi
AUC
Cost
Fi
AUC
Cost
Strategy
LR RF
0.539
0.935
$290
0.452
0.797
$591
0.506
0.825
$591
Full
0.605
0.938
$290
0.439
0.764
$591
0.456
0.801
$591
Full
XGBoost
0.617
0.938
$290
0.404
0.785
$591
0.431
0.828
$591
Full
LightGBM
0.627
0.941
$290
0.474
0.790
$591
0.500
0.844
$591
Full
3-layer DNN
0.616
0.938
$290
0.494
0.802
$591
0.517
0.845
$591
Full
LR (2 panels)
0,401
0.859
$92
0.473
0.797
S92
0.488
0.811
S92
Fixed
RF punels)
0.504
0.887
$92
0.425
0.768
S92
0.478
0.828
$92
Fixed
XGBoost (2 pancls)
0,519
0.895
$92
0,410
0.781
S92
0.459
0.877
S92
Fixed
LightGBM (2 panels)
0.571
0,901
S92
0,491
0.792
S92
0.502
0.864
S92
Fixed
FS
0.585
0.927
S74
0.434
0.787
S98
0,500
0.837
$90
Fixed
RS
0.437
0.845
$145
0.424
0.748
$295
0.473
0.789
$295
Random
CWCF
0.554
0.718
$256
0.283
0,510
$326
0.112
0.503
$301
Dynamic
SM-DDPO, pretrainod
0.607
0.925
S80
0.519
0.789
S90
0.567
0.836
S85
Dynamic
SM-DDPO snd2end
0.624
0.928
S62
0.495
0.795
597
0.562
0.845
590
Dynamic
Comparisons with baseline models using full observation of data. The results are presented in Table 2 Across all three clinical tasks, our proposed model can achieve comparable or even state-of-the-art performance, while significantly reducing the financial cost. On sepsis dataset, SM-DDPO yielded better results (F,=0.563 AUROC=0.845 than the strongest baseline mod- els, LightGBM =F,70,517, AUROC=0.845 when saving up 84% in test cost. On ferritin dataset, LightGBM (F,20,627, AUROC=0,948) performed slightly better than our model (F,=0.624, AU- ROC=0.928), however, by using 5x testing cost. On AKI dataset, RM-DDPOendZenn (F)=0.495, AUROC=0. 795) achieved comparable results to the optimal full observation model, 3-layer MLP (F)=0.494, A.80C= while saving the testing cost from $591 $90.
Comparisons with other test selection strategies. Our proposed SM-DDPOcnt2end using RL- inspired dynamic selection strategy, consistently yielded better performance and required less testing cost across all three datasets, when compared to the models using fixed or random selection strategy. For fixed test selection strategy, we first tested the classification methods using the two most relevant panels: CBC and CMP. These baselines with reduced testing cost still behaved much worse than our approach in both F1 score and AUROC. We also tested another fixed selection (FS) baseline, where chose to always observe 2 most selected test panels reported in our approach for all patients, while keeping other modules the same. Our approach outperformed FS on both Fi score and AUROC while having similar testing cost. The random selection (RS) baseline selected test panels uniformly at random and had a worse performance. Q-learning for classification with costly Janish) this is because the model uses the same network for selecting tests and learning rewards. For such imbalanced datasets, this may make the training unstable and difficult to optimize.
features (CWCF) Janisch et al. (2019) performed poorly on all three clinical datasets. We believe
Efficiency and accuracy of end-to-end training. The classifier and panel selector of RM-DDPOcsd2ond are both trained from the scratch, using Algorithm [ As in Table 2 this end-to-end training scheme gives comparable accuracy with policies that rely on a heavily pretrained classifier (SM-DD)POpnetran)) If brand-new diagnostic/predictiee task is given, our algorithm can be trained without prior data/knowledge about the new disease. It can adaptively prescribe lab tests to learn the disease model and test selection policy in an online fashion. End-to-end training is also more data-efficient and runtime efficient.
Interpretability. Our algorithm is able to select test panels that are clinically relevant. For ferritin prediction, our algorithm identifies TSAT as most important panel, which is indeed useful for detecting iron deficiency. For AKI prediction, our algorithm recommends serum creatinine level test as important predictor for 95% of subjects, i.c.. current and past serum creatinine is indicative of future AKI, expanding its utility as a biomarker de Geus et al. (2012).
8

---

Published as a conference paper at ICLR 2023
6.3 TRAINING CURVES
We present the training curves on AKI dataset in Figure We refer more results to Appendix
SM-DDPO learns the disease model. In end-to-end training, the diagnostic classifier is trained from the scratch. It maps any partially-observed patient state to a diagnosis/rredittion We evaluate this classifier static data distributions, in order eliminate the effect of dynamic test selection and focus on classification quality. Figure shows that the classifier learns to make the high-quality prediction with improved quality during RL, via training only on data selected by the RL algorithm.
0.70
Cross Entropy Lozs
rardom missing pattern
0.45
RL missing pettern
0.7
Score
0.40
AUROC
0.65
0.35
0.6
0.60
0.30
0.5
0
50 100
0
50
100
0
50
100
End-to-end training loops
End-ta-end training loops
End-to-end training loops
Figure 3: Classifier improvement during RL. training on AKI Dataset. Accuracy of the learned classifier is evaluated on static patient distributions, with 1) random missing pattern, where we uniformly at random augment the test data; 2) missing pattern of the optimal policy's ss state distribution. During the end-to-end RL training. the classifier gradually improves and has higher accuracy with the second missing pattern.
Ferritin Dataset
AKI Dataset
Sepsis Dataset
0.6
0.4
c: sccore
0.4
sccor
0.4
0..2
k 0.2
0.2
Ugr
0.0
0.0
0.0
100
200
300
200
400
600
200
400
600
Testing Cost
Testing Cost
Testing Cost
0.8
0.9
0.8
0.7
AUROC
AUROC
0.8
0.6
0.7
0.5
0.6
Cope ConC
0.7
Dos
0.4
0.5
0
100
200
300
200
400
600
200
400
600
Testing Cost
Testing Cost
Testing Cost
Figure 4: Cost-Fi Pareto Front for maximizing Fi Score on Ferritin, AKI and Sepsis Datasets
6.4 COST-F, PARETO FRONT
Figure 4 illustrates the Pareto fronts learned on all three datasets. We trained for optimal policies on 190 MDP instances specified by different value pairs of (A, p) in Theorem 4.1] and present the corresponding performance on Fj score (red) and AUROC (blue) evaluated on the test sets. We identify the Pareto front the upper envelope of these solutions, which are the yellow curves in Figure 4 These results present the full tradeoff between testing cost and diagnostic/predietiee accuracy. As corollary, given any cost budget B, one is able to obtain the best testing strategy with the optimal F1 performance directly from Figure 4 We present a zoom-in version in Appendix (B]
SUMMARY
In this work, we develop Semi-Model based Deep Diagnosis Policy Optimization (SM-DDPO) method to find optimal cost-sensitive dynamic policies and achieve state-of-art performances on real-world clinical datasets with up to 85% reduction in testing cost.
9

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
Mengdi Wang acknowledges the support by NSF grants DMS-195360 IIS-2107304, CMMI- 1653435, ONR grant 1006977, and http://C3.AI.
Yuan Luo acknowledges the support by NIH grants UOITR003528 and R01LM013337.
Yikuan Li acknowledges the support by AHA grant 23PRE1010660.
REFERENCES
Derek C Angus and Tom Van der Poll. Severe sepsis and septic shock. N Engl Med, 369:840-851, 2013.
Peter C Austin, Ian R White, Douglas S Lee, and Stef van Buuren. Missing data in clinical research: a tutorial on multiple imputation. Canadian Journal of Cardiology, 37(9):132--133) 2021.
Jinbo Bi, Kristin Bennett, Mark Embrechts, Curt Breneman, and Minghu Song. Dimensionality reduction via sparse support vector machines. Journal of Machine Learning Research, 3(Mar): 1229-1243, 2003.
Mustafa Bilgic and Lise Getoor. Voila: Efficient feature-value acquisition for classification. In Proceedings of the national conference on artificial intelligence. volume 22, pp. 1225. Menlo Park, CA; Cambridge, MA; London; AAAI Press; MIT Press; 1999, 2007.
Xiaoyong Chai, Lin Deng, Qiang Yang, and Charles X Ling. Test-cost sensitive naive bayes classification, In Fourth IEEE International Conference on Data Mining (ICDM 04), pp. 51-58. IEEE, 2004.
Tianqi Chen and Carlos Guestrin. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd acm sigkdd international conference on knowledge discovery and data mining, pp. 785-794, 2016.
Gabriella Contardo, Ludovic Denoyer, and Thierry Artières. Sequential cost-sensitive feature acquisition. In International symposium on intelligent data analysis, pp. 284-294. Springer, 2016.
Brian K Crownover and Carlton Covey. Hereditary hemochromatosis. American family physician, 87 (3):183-190, 2013.
Hilde RH de Geus, Michiel G Betjes, and Jan Bakker. Biomarkers for the prediction of acute kidney injury: a narrative review on current status and future challenges. Clinical kidney journal, 5(2): 102-108, 2012.
Menghan Ding and Yuan Luo. Unsupervised phenotyping of sepsis using nonnegative matrix factorization of temporal trends from a multivariate panel of physiological measurements. BMC Medical Informatics and Decision Making, 21(5):1-15, 2021.
Conor Fahy and Shengxiang Yang. Dynamic feature selection for clustering high dimensional data streams. IEEE Access, 1127128-121440 2019.
Omer Gottesman, Fredrik Johansson, Matthieu Komorowski, Aldo Faisal, David Sontag, Finale Doshi-Veles and Leo Anthony Celi. Guidelines for reinforcement learning healtheare. Nature medicine, 25(1):16-188 2019.
Arthur Guez, Robert D Vincent, Massimo Avoli, and Joelle Pineau. Adaptive treatment of epilepsy batch-mode reinforcement learning. In AAAI, pp. 1671-1678, 2008
Isabelle Guyon and André Elisseeff. An introduction to variable and feature selection. Journal of 2003. 2003.
machine learning research,
He He. Hal Daumé III. and Jason Eisner. Cost-sensitive dynamic feature selection. In ICML Inferning Workshop, 2012.
Tin Kam Ho. Random decision forests. In Proceedings of 3rd international conference on document analysis and recognition, volume pp. 278-282. IEEF, 1995.
10

---

Published as a conference paper at ICLR 2023
Amelia Huck and Kent Lewandrowski. Utilization management in the clinical laboratory: introduction and overview of the literature. Clinica Chimica Acta, 427:111-117, 2014.
Jaromír Janisch, Tomá Pevný, and Viliam Lisý. Classification with costly features using deep rein- forcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 3959-3966, 2019.
Shihao Ji and Lawrence Carin. Cost-sensitive feature acquisition and classification. Puttern Recogni- lion, 4065):144-14485 2007.
Alistair EW Johnson, Tom J Pollard, Lu Shen, Li-wei H Lehman, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G Mark. Mimic-iii, freely accessible critical care database. Scientific data, 3(1):1-9, 2016.
Sergey Karayev, Mario Fritz, and Trevor Darrell. Dynamic feature selection for classification on budget. In International conference on machine learning (ICML): Workshop on prediction with sequential models. Citeseer, 2013.
Guolin Kc, Qi Meng, Thomas Finley, Taifeng Wang. Wei Chen, Weidong Ma, Qiwei Ye, and Tie-Yan Liu. Lightghm: highly efficient gradient boosting decision tree. Advances neural information processing systems, 30, 2017.
John A Kellum and Norbert Lameire. Diagnosis, evaluation, and management acute kidney injury: kdigo summary (part 1). Critical care, 17(1):204, 2013.
Matloob Khushi, Kamran Shaukat, Talha Mahboob Alam, Ibrahim A Hameed, Shahadat Uddin,
Suhuai Luo, Xiaoyan Yang, and Maranatha Consuelo Reyes. A comparative performance analysis of data resampling methods imbalance medical data. IEEE Access, 9: 9:10990--099755 2021.
Ron Kohavi and George H John. Wrappers for feature subset selection. Artificial intelligence, 97 (1-2):273-324, 1997.
Matthieu Komorowski, Leo A Celi, Omar Badawi, Anthony Gordon, and A Aldo Faisal. The artificial intelligence clinician learns optimal treatment strategies for sepsis in intensive care. Nature medicine, 24(11):
2018.
Steef Kurstjens, Thomas de Bel, Armando van der Horst, Ron Kusters, Johannes Krabbe, and Jasmijn van Balveren. Automated prediction of low ferritin concentrations using machine learning algorithm. Clinical Chemistry and Laboratory Medicine (CCLM), 2022.
Eun Hye Lee, Mi Hwa Shin, Jong-Min Park, Sang-Guk. Lee, Nam Su Ku, Young Sam Kim, Moo Suk Park, Jae-Chul Pyun, and Kyung Soo Chung. Diagnosis and mortality prediction of sepsis via lysophosphatidylcholine 16: 0 measured by maldi-tof ms. Scientific reports, 10(1):1-9, 2020.
Der-Chiang Li, Chiao- Wen Liu, and Susan C Hu. A learning method for the class imbalance problem with medical data sets. Computers in biology and medicine, 40(5):509-515 2010.
Yang Li and Junier Oliva. Active feature acquisition with generative surrogate models. In International Conference on Machine Learning, pp. 6450-6451 PMLR, 2021.
Wei-Chao Lin and Chih-Fong Tsai. Missing value imputation: review and analysis of the literature (2006-2017). Artificial Intelligence Review, 53(2): 1487-1599, 2020.
Yuan Luo, Peter Szolovits, Anand S Dighe, and Jason M Baron. Using machine learning to predict laboratory test results. American journal of clinical pathology, 145(6): 2016.
Qi Ma and Sujit K Ghosh. Emflow: Data imputation in latent space via em and deep flow models. arXiv preprint arXiv:2106.0488 2021.
Aditya Menon, Harikrishna Narasimhan, Shivani Agarwal, and Sanjay Chawla. On the statistical consistency of algorithms for binary classification under class imbalance. In Sanjoy Dasgupta and David McAllester (eds.), Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pp. 603 611 Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR. URL https: / http://///ccceeinncs.. mlr. press / v28 / 2722/mmnnnn3a html

---

Published as a conference paper at ICLR 2023
Rui P Moreno, Barbara Metnitz, Leopold Adler, Anette Hoechtl, Peter Bauer, and Philipp GH Metnitz. Sepsis mortality prediction based on predisposition, infection and response. Intensive care medicine, 34(3):496-555 2008.
HyunJi Alex Nam, Scott Fleming, and Emma Brunskill. Reinforcement learning with state observation costs in action-contingent noiselessly observable markov decision processes. Advances in Neural Information Processing Systems, 34:15650-15666, 2021.
Nagarajan Natarajan, Inderjit S.
Dhillon, Pradeep Ravikumar, and Ambuj Tewari. Cost-sensitive learning with noisy labels. Journal of Machine Learning Research, 18(155):1-33, 2018. URL
http:// /ttt://jmlr.... http:///ggggapp.///////s//////// htm
Muhammad S Osman, Adnan M Abu-Mahfouz, and Philip R Page. A survey on data imputation techniques: Water distribution system as a use case. IEEE Access, 6:63279 2x63279-63291, 2018.
Mark A Perazella. The urine sediment as a biomarker of kidney disease. American Journal of Kidney Diseases, 66(5):748- 755,, 2015.
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines:: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268): 2021. URL pnttp:///mml org patpg//ppp///22/// 20 1364.htmm
M Mostafizur Rahman and Darryl N Davis. Addressing the class imbalance problem in medical datasets. International Journal of Machine Learning and Computing, 3(2):224, 2013.
Nelson Sanchez-Pinto and Robinder G Khemani. Development of prediction model of early acute kidney injury in critically ill children using electronic health record data. Pediatric Critical Care Medicine, 17(6))508 -1588555117100 2016.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839) 609, 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: arXiv: 1707.06347, 2017.
Jiyoung Shin, Yikuan Li, and Yuan Luo. Early prediction of mortality in critical care setting in sepsis patients using structured features and unstructured clinical notes. In 202/ IEEE International Conference on Bioinformatics and Biomedicine (BIBM), pp. 2885-2890. IEEE, 2021.
Matthew W Short and Jason E Domagalski. Iron deficiency anemia: evaluation and management. American family physician, 87(2)988-044 2013.
Mervyn Singer, Clifford S Deutschman, Christopher Warren Seymour, Manu Shankar-Har Djillali Annane, Michael Bauer, Rinaldo Bellomo, Gordon R Bernard, Jean-Daniel Chiche, Craig M Coopersmith, et al. The third international consensus definitions for sepsis and septic shock (sepsis-3). Jama, 315(8)8001-80,, 2016.
Mengxin Sun, Jason Baron, Anand Dighe, Peter Szolovits, Richard G Wunderink, Tamara Isakova, and Yuan Luo. Early prediction of acute kidney injury in critical care setting using clinical notes and structured multivariate physiological measurements. MedInfo, 364:368-372. 2019.
Kirill Trapeznikov and Venkatesh Saligrama. Supervised sequential classification under budget constraints. In Artificial intelligence and statistics, pp. 581-589, PMLR, 2013.
Jason Weston, Sayan Mukherjee, Olivier Chapelle, Massimiliano Pontil, Tomaso Poggio, and Vladimir Vapnik. Feature selection for svms. Advances in neural information processing systems, 13, 2000.
Jason Weston, André Elisseeff, Bernhard Schölkopf. and Mike Tipping. Use of the zero with linear models and kernel methods. The Journal of Machine Learning Research, 2:1439-14611 2003.
12

---

Published as a conference paper at ICLR 2023
Haiyan Yin, Yingzhen Li, Sinno Jialin Pan, Cheng Zhang, and Sebastian Tschiatschek. Reinforcement learning with efficient active feature acquisition. arXiv preprint arXiv:2011.00825, 2020.
Chao Yu, Jiming Liu, Shamim Nemati, and Guosheng Yin. Reinforcement learning in healthcare: A survey, ACM Computing Surveys (CSUR), 55(1):1-36, 2021.
Junyu Zhang, Alec Koppel, Amrit Singh Bedi, Csaba Szepesvari, and Mengdi Wang. Variational policy gradient method for reinforcement learning with general utilities. Advances in Neural Information Processing Systems, 33:4572 4583, 2020.
Lindsay P Zimmerman, Paul A Reyfman, Angela DR Smith, Zexian Zeng, Abel Kho, L Nelson Manchez-PPinto, and Yuan Luo. Early prediction of acute kidney injury following icu admission using a multivariate panel of physiological measurements, BMC medical informatics and decision making, 19(1):1-12, 2019.
13

---

Published as a conference paper at ICLR 2023
APPENDIX
We organized the appendices as follows. We first provide a complete literature review in appendix A In appendix BJ we provided the descriptive statistics of all three clinical datasets, followed by the experimental results and detailed training settings. In appendix [ we described the EMFlow imputation algorithm. In appendix D we gave more details of the proposed end-to-end training approach. In appendix [ we proved the Theorem 4.1, In appendix F we discuss the alternative AM metric for evaluating the diagnostic performance as an example to show the capability of our framework to handle classic linear metrics.
A RELATED WORK
Reinforcement learning (RL) has been applied in multiple clinical care settings to learn optimal treatment strategies for sepsis Komorowski al. (2018). to customize antiepilepsy drugs for seizure control Guez et al. (2088) etc. See survey Yu et al. (2021) for more comprehensive summary. Guidelines on using RL for optimizing treatments in healthcare has also been proposed around the topics of variable availability, sample size for policy evaluation, and how to ensure learned policy works prospectively as intended Gottesman et al. (2019). However, using RL for simultaneously reducing the healthcare cost and improving patient's outcomes has been underexplored.
Our problem of cost-sensitive dynamic diagnosis/prediction is closely related to feature selection in supervised learning. The original static feature selection methods, where there exists common subset of features selected for all inputs, were extensively discussed Guyon & Elisseeff (2003). including filter models (e.g., variable ranking), wrapper approaches Kohavi & John (1997), ), well as integration of the feature selection in the learning process by using &l fl-norm and 20-norm Bi et al. (2003); Weston et al. (2003) 2000). Dynamic feature selection methods He et al. (2012). also named as adaptive feature acquisition Contardo et al. (2016) or feature selection with budget Karayev et al. (2013). were then proposed to take the difference between inputs into account. Different subsets of features are selected with respect to different inputs. By defining certain information value of the features Fahy Yang (2019); Bilgic & Getoor (007).. or estimating the gain of acquiring new feature would yield Chai et al. (2004). these methods are able to select appropriate features in an adaptive way to balance the costs of the selected features and the learning quality.
Recently, reinforcement learning based approaches Ji & Carin (2077); y Trapeznikov & Saligrama :013):Danisch Danisch al. (2019); Yin et al. (2020): Li & to dynamically select features for peeeiition/classiiiaatt or pre-define a sum of rewards as objectives. While our task of deriving the cost-Fi F1 Pareto front, maximizing F1 score under cost budget constraints, is much harder and cannot be ex expsseed using pre-defined reward functions. Work Nam et al.] (202T); I): Jī & Carin (2007): Yin et al. (2020) models the feature selection process as a partially observed Markov Decision Process (POMDP).. &
(2013);
Nam et al. (2021) are also proposed All these related works are directly given,
Carin (007) points out that POMDP solutions are hard to train and cannot handle continuous-valuse observations. A similar approach is proposed byJansshh al. (019)) using deep Q learning method. In our approach, we use modular network and a combination of various training schemes for more efficient RL. TTrapznnkkvv & Saligrama (2013) ranks features in a prefixed order and only requires the agent to decide if it will add a next feature or stop. The prefixed ranked order makes the problem easier to solve, but it is suboptimal compared to fully dynamic policies.
In this work, we deal with datasets with highly imbalancement, i.e., there are much fewer ill patients compared to healthy patients. To handle such data imbalancement, we use the metric of Fl score, which is a well-established and acknowledged evaluation metric in the medical domain Perazella (2015); SSanchez-Pinto
Khemani (2016), to evaluate the diagnostic performance. Alternative metrics includes linear AM metries Natarajan al. (2018); Menon et al. (2013) for imbalanced datasets can also handled by our frameworks, which we discussed in Appendix
B EXPERIMENTS
B. DESCRIPTION OF DATASETS
We summarize a detailed description of all lab tests in three datasets in Table 3 4 and 5
14

---

Published as a conference paper at ICLR 2023
Table 3: Summary descriptive statisties of the Ferritin dataset. (N 43472. Binary variables are reported with positive numbers and percentages. Continuous variables are reported with means and interquartile ranges.
Variables
meun % S/QQ
Variables
mean
/IQR
Demographics
Complete Blood Count Panel ($44)
Age
(37.00,
'Abs. Nacleated RBC Count (K/yL)
0.21
(0.00, 0.00)
Race/bchnicity
Abs. Basophils (K/oL.)
0.05
(0.00, 0.10)
Asian
26887
5,18%
Abs. Bosinophils (K/₂L)
0.16
(0.10, 0.20)
Black
14.61%
Abs. Lymphocytes (K/yl.)
2,05
(1.45,
Hispanic
4272
9.83%
Absolute Monocytes (K/jL)
0.53
(0.40, 0.60)
White
3708
61.85%
Absolute Neutrophils (K/pil.)
4,13
(2,90,
Other
2251
8,53%
Percent Eosinophils (5K)
2,46
(1.00,
Comprehensive Metabolic Panel (S448)
Percent Lymphooytes (56
29.49
(23.67, 35.00)
Alanine Transuminase (UL)
20.14
(12.00, 21
Percent Monocytes (5)
8
16.00, 9.00)
Albumin (g/0L)
4.3
(4.10, 4,50)
Percent Neutrophils (5)
60.23
(53.50, 67.00)
Alkaline Phosphutase (C/L.)
67.03
(48.00) 77.00)
Percent Rasophils (9)
0.68
(0.00), 1.00
Anion Gap (mEE))))
9.79
(8.00, 12.00)
"hhite Blaod Count (K(jL)
6.92
(5.30, 7.90)
Aspurtate Trunsaminase (UVL.)
22,06
(16.00, 23.00)
"Red Blood Cell Count (M/jl.)
4.34
(4.07, 4.67)
"Bicarborste (mmoltl.)
27.09
(26.00, 29:00)
Plittelet (K/pl.)
274.78
(224.00, 50)
"Blood Urea Nitrogen (mg/lL)
15.24
(11,00, 17.00)
Mcan cell volume (IL)
28.91
(27.70, 30.70)
Calciiu (mp/dL))
9.49
(9.20, 9.75)
Mcan cell hemoglobin (pg/RBC)
32.07
(81,38, 32,95)
Chloride (mmol/.)
103.71
(102.00, 105.00)
Mean cell hemoglobin concentration (pkiL)
90.03
(87.00 94.00)
Creatiiinee (mmok/))
0.87
(0.69. 0.89)
RBC distribution width (%)
13.89
(12.50, 14.60)
Sooiium (mmol/L)
139,27
(138.00, 141.00)
Hematocrit (9%
38.87
(36.60) 41,80)
Tocal Bilirubin (mg/ll.)
0.55
(0.40, 0.60)
Hemoglobin (g/llL)
12.48
(11.60, 13.60)
Total Protein (gidl.)
7.22
(6.90, 7.50)
Transferrin Saturation Test (540)
Potassiom (mmmol/L)
4.11
(3,90, 4.30)
Iron (pg/dL))
82.33
(54.00, 104.00)
*Glncose (mg/L)
99.88
(86,00, 102.00)
Total Iron- Iron-binding Capacity (ug/lLL)
366.81
(321.00. 409.00)
Globalin (gldL)
3.19
(2.70, 3.60)
Outcome
Vitamin B-12 Test (866)
Abnormal Ferritin Level
4649
10.69%
Cobalamin (pg/tl))
$47.8
(302.00, 668.00)
These variables can also be ordered through basic blood count panel ($26). These variables can also he ordered through basic metabolic panel ($36).
Table 4: Summary descriptive statistics of the AKI dataset. (N 23950. Binary variables are reported with positive numbers and percentages. Continuous variables are reported with means and interquartile ranges.)
Varlables
N/Mcan
%/TQR
Varlables
KiMcan
FETQQ
Demographies
Comprehensive Melabolic Pane (548)
Gender
Glucose Lovel Maximum (mg/dL)
173.36
(128.00, 196.00)
Fomale
9755
40.73%
Bicarbonatc Level Minimum (mp/il.)
23,87
(21,00,26 26.00)
Mule
14195
59.27%
Cremtnine Level Minimum (mgAll))
0.74
(0.66, 0.90)
Age (yr)
60.89
(50.66, 73 73.50)
Crestinine Levll Maximom (mg)il))
0.8
(0.60) 1,00)
Rawe/iiinicity
Utea Nirugen Level Maximum (mghlL)
(11.00,20 20.000
Black
1688
Calcium Level Minimums (mpill.)
8.08
(7.62, 8.54)
Hispanic
845
3.53%
Estimated Giomeralar Pilraatinn Rote (eGPR)
110.2
(80.94, 124. 30)
White
17261
72.07%
Potassium Level Maximum (mg/d..)
4.39
(3:90, 4.70)
Other
44%
Activated Partial Thromboplastin Time Pand
Vital Signs Bedside Mensurements
Partial Thrembopliatin Time Maximom (4)
(36.60, 34.60)
Heart Rate Maximum (bpm)
105.4
(91.00, 117.00)
Partial Thromboplastin Time Minimum (6)
40.22
(27.80,41
Heart Rate Menn (bpm)
86.97
(76.58, 96. 96)
International Normalized Ratio Minimum
1.34
(1.10, 1.40)
Systalic Minimum (mmigg)
92.51
102.00)
Interational Nomnalized Ratio Maximum
1.48
(1.20) 1.52)
Systolic Meon (mmilg)
118.87
(107.34, 458.45)
Prochrombin Time Minimum
(13.10, 15.08)
Diastolic BP Minimam (mmllg)
45.31
(39.00, 52.00)
Prochrombin Time Maximum (5)
15.63
(13,40, 16.18)
Diastelic BP Mean (mmilg)
61.98
(51.88, 67.87)
Complete Blood Count Panel (544)
Mechanical Ventilation
12273
51.24%
Hemoglobin Level Minimum (g/lL) Platelet Count
10.33
(8.90, 11.70)
Average Output
2202.65
2792.00)
210.24
(134.00 261.066000 (5)00)
Temperature Maximam (celiuus)
37.6
(37.06, 38.06)
White Blood Cell Count Maximum (mg/lL))
12.78
(8.50, 15.50)
Arlecial Blood Gas Test ($473)
Outcome
Sp02 Minimum (5)
92.23
(91.00, 95.
Acute Kidney Injury
3945
16.47%
Sp02 Maximum
98.
B.2 SUPPLEMENT EXPERIMENT RESULTS
In this section, we show a more detailed results of Table presented in the main text.
2
B.3 DETAILED TRAINING SETTINGS
The experiments are conducted in part through the computational resources and staff contributions provided for the Quest high performance computing facility at Northwestern University which is jointly supported by the Office of the Provost, the Office for Research, and Northwestern University Information Technology. Quest provides computing access of over 11,800 CPU cores. In our experiment, we deploy each model training job one CPU, so that multiple
Computing Resource
15

---

Published as a conference paper at ICLR 2023
Table 5: Summary descriptive statistics of the Sepsis dataset. (N 5396. Binary variables are reported with positive numbers and percentages. Continuous variables are reported with means and interquartile ranges.)
Varlables
N/Mean % IQR
Variables
N/Mean
%/IQR
Demographics
Vital Signs, Bedside Meusurements, Medical Histories
Geader
Heart Rate (bpm)
90.99
(77.00, 104,00)
Female
3014
55 55.86%
Respiration Rate (insp/min)
19.37
(15.00, 23,00)
Male
2382
Systolic (mmFlg)
124.46
(107.00, 141.00)
Age (ye)
65.5
(53.90, 79.93)
Diastolic BP (mmilg)
66.75
(55.00, 77,00)
RacerEitnicitt
Glasgow Comma Scale
4.91
6.00)
Asian
167
3 3.09%
Mcchanical Ventilation
2781
48. 11%
Black
473
8.77%
Urine Ourput (mll
214.91
(70.00, 300.00)
Ilsspanic
182
3.37%
Tidal Volume Temperature (celsius)
496, 98
(450.00. 550.00)
White
3923
72.70%
36,64
(36.10, 37.22)
Other
654
12.12%
History Metastatic Cancer
342
5.92%
Marital Status
History of Diabetes
1629
28, 18%
Divorced
328
6.08%
28.54
(23.53, 31.75)
Married
2379
44,095%
Comprehensive Metabolic Panel ($48)
Single
28.47%
Glucose Level (mg/dL)
4.91
5.11)
Widowed
795
14,75 43.%
Bicarbonate Lovel (mEq/LL
22.95
(20.00, 26.00)
Unknown
358
6.63%
Scrum Creatinine (meldl.)
1.54
(0.80, 1.60)
Admission Type
Aspurtate uninoorramserasee (IU/L)
4,01
(3.18, 4.49)
Elective
310
5.74%
Bilirubin (mg/dL)
1,57
(0.40, 1,30)
Emergency
5026
93.14%
Chloride (mEg/L)
105.23
(101.00, 00)
Urgent
60
1.11%
Carbon Dioxide (mFFq/1.)
24.61
(21 10,00, 28,00)
Insurance Type
Lactate (ammol/1.)
2.16
(1.10, 2.60)
Government
157
2.91%
Magnesium (mg/dl))
1.92
(1.70, 2. 10)
Medicaid
532
9. 9.86%
Blood Urca Nitrogen Lovel (mng/dL)
28.85
35.00)
Medicare
3123
57 88%
Blood albumin (g/dl.)
3.01
(2.60, 3, 3.40)
Private
1538
28.50%
Blood sodium (mEEq)..)
138.3
(135.00. 141.00)
Self-pay
46
0.85%
Potassium Level (mEyL))
4.18
(3.70, 4.60)
Arterial Blood Gas Test (8473)
Complete Blood Count Panel (544)
Base exeess (mFq/L))
-1.51
(-4.00, 1.00)
Hemoglobin Level (pidL)
10.83
(9.40, 12.20)
PH (unit)
7.36
(7.31, 7.43)
Hematocrit (%)
32.37
(28. 30, 36.30)
Fraction of Inspired Oxygen (S)
72.01
(50.00, 100.00)
Platelet
(K/pL)
5.21
(4.93, 5.59)
Mean Arterial Pressure (mm)Hg)
82.08
(69.00, 93,00)
Red Blood Coll Count (K/xL)
3.58
(3.10, 4.05)
Blood Oxygen Saturation (%)
96.73
(95.00, 100.00
White Blood Cell Count (K/jL)
12.69
(7.50, 15,30)
Composite Score
Activated Partial Thromboplastin Time Test Panel ($26)
Sequential Organ Failuure Assessment
4.63
(2.00, 6.00)
Partial Thromboplastin Time (s)
36.75
(26.50, 37.30)
Outcome
International Normalized Ratio
1.5
(1.10, 1.50)
In-bospital Mortality
3945
16.47%
configurations can be tested simultancously. Each training job requires a wall-time less than 2 hours of single CPU core.
Hyper- Parree-Parrseter Settings The hyper-parameter of our model is divided onto four parts: parame- ter for training imputer, classifier, PPO policy, well the end-to-end training parameters. We list all the hyper-parameters we tuned in Table 7 including both the tuning range and final selection. The unmentioned parameters training EMFlow imputer are set to the same as its original work Ma Ghosh (2021). The unmentioned parameters in PPO are set to default values in Python package Raffin et al. 2021).
Data Splitting Scheme We split each dataset into 3 parts: training set (75%), validation set (15%), and test set (10%). Training data is further splitted into two disjoint sets, one for pretraining the state encoder (25%) and the other for end-to-end RL training (50%). The validation set is splitted in the same way, one for tuning parammtrree for training the state encoder (5%) and the other one for tuning syper-parameters for end-to-end RL training (10%). Here, all percentages calculated w.r.t. the whole dataset size.
Policy Improvement During RL Training In Figure we showed the performance of the RL policies evaluated on train, validation and test sets. We note that the three curves closely match one another, confirming the generalizability of the learned dynamic classification policy.
Zoom-in Version of the cost-F; Pareto Front Here we present the cost-F1 Pareto front, comparing only to baselines with similar testing cost (Figure 6 is the envelope of the 190 solutions, is the final Pareto front we obtained the test set. Comparing to baselines with similar cost, our approach gives a better performance on F, score under the same cost budget. As already shown in Figure 4 even comparing to fully observable approaches which have
is a subset of Figure 4 The yellow curve, which
16

---

Published as a conference paper at ICLR 2023
Table Complete version of Table 2 Model performance, measured by score, area under ROC (AUC), and testing cost, for three real-world clinical datasets. The tested models include logistic regression (LR), random forests (RF), gradient boosted regression trees (XGBoost Chen Guestrin (2016) and LightGBM Ke al. (2017) ). 3-layer multi-layer perceptron, Q-learning for classification with costly features (CWCF) Janisch et al. (2019), Random Selection (RS), Fixed Selection (FS). Among them, classical methods including LR, RF, XGBoost, Lightbgm are tested under both fully observable case and fixed observable case where most relavant panels (CMB and CBCP panel) chosen. All models were fine-tuned to maximize the Fi score. The model yielded the highest Fj score is bold. The model required the least testing cost is underlined.
Dataset
Models
Test Selection Strategy
Fi
AUC
Cost
Ferritin
LR RF XGBoost
Full
0,541 + 0.010
0.934 + 0.008
S290
Full
0.597 : 0.024
0.931 +: 0.006
S290
Full
0.610 + 0.027
0.940 + 0.004
$290
Lightbgm
Full
0.625 + 0.011
0.941 + 0.006
$290
3-layer DNN
Full
0.615 = 0.018
0.937 + 0.005
$290
LR (2 panels)
Fixed
0,401 + 0.016
0.859 + 0.013
$92
RF (2 panels)
Fixed
0.504 + 0.023
0.887 + 0.012
S92
XGBoost (2 panels)
Fixed
0.519 0.015
0.895 0.007 0.007
$92
Lightbgm (2 panels)
Fixed
0.571 = 0.010
0.901 = 0.006
$92
FS
Fixed
0.588 = 0.017
0.922 0.007
S74
RS
Random
0,427 = 0.023
0.837 + 0.012
S144.78 0.52
CWCF
Dynamic
0.531 + 0.021
0.702 + 0.018
S261.02 + 1.20
SM-DDPO pretrained
Dynamic
0.587 - 0.021
0.922 0.008
S80.72 + 0.63
SM-DDPO endlend
Dynamic
0.613 - 0.018
0.926 0.006
$62.51 + 0.46
AKI
LR
Full
0.453 = 0.028
0.791 + 0.018
S591
RF
Full
0.438 = 0.032
0.771 + 0.016
S591
XGBoost
Full
0.414 + 0.023
0.789 + 0.017
S591
Lighthgm
Full
0.474 + 0.017
0.791 + 0.012
$591
3-layer DNN
Full
0.492 + 0.024
0.797 + 0.010
S591
LR (2 panels)
Fixed
0.473 + 0.006
0.797 + 0.005
$92
RF (2 panels)
Fixed
0,425 + 0,016
0.768 + 0.004
S92
XGBoost (2 panels)
Fixed
0.410 = 0.022
0.781 = 0.009
S92
Lighthgm (2 panels)
Fixed
0.491 = 0.011
0.792 + 0.006
S92
FS
Fixed
0.432 + 0.028
0.789 + 0.015
S98
RS
Random
0.411 + 0.029
0.744 0.019
$294.88
1.32
CWCF
Dynamic
0.278 + 0.053
0.510 + 0.009
$340.39 + 19.01
SM-DDPO peetrained
Dynamic
0.514 + 0.030
0.788 + 0.022
$90.12 + 0.68
SM-DDPO endlend
Dynamic
0.489 + 0.026
0.796 + 0.019
$97.31 + 0.58
Sepsis
LR
Full
0.507 = 0.035
0.814 + 0.030
$591
RF
Full
0.442 = 0.051
0.820 + 0.016
$591
XGBoost
Full
0.451 = 0.051
0.842 + 0.025
$591
Lightbgm
Full
0.501 = 0.046
0.850 + 0.018
$591
3-layer DNN
Full
0.515 + 0.031
0.842 + 0.014
$591
LR (2 panels)
Fixed
0.488 + 0.021
0,811 + 0.015
$92
RF (2 panels)
Fixed
0.478 + 0.053
0.828 + 0.018
$92
XGBoust (2 panels)
Fixed
0.459 + 0.032
0.877 + 0.036
$92
Lightbgm (2 panels)
Fixed
0.502 = 0.031
0.864 + 0.010
$92
FS
Fixed
0.489 + 0.039
0.830 + 0.020
$90
RS
Random
0.475 + 0.062
0.779 + 0.030
$295.01 + 2.41
CWCF
Dynamic
0.118 + 0.011
0.503 + 0.001
$401.83
120.76
SM-DDPO (eetrained
Dynamic
0.545 + 0.034
0.834 0.022
$86.21
144
SM-DDPO...th endlend
Dynamic
0.532 - 0.036
0.841 + 0.023
$90.55 = 1.20
much higher testing cost, approach exhibits comparing performance on Fi score and AUROC with a much lower testing cost.
Lastly, emphasize that approach gives the whole Pareto front, which represents complete and rigorous characterization of the accuracy-cos: trade-off in our medical diagnostics tasks. While all the baselines compared only gives single points this trade-off characterization.
17

---

Published as a conference paper at ICLR 2023
Table 7: Hyper-parameter Table
Hyper-parameters
Ferritin
AKI
Sepsis
Selection Range
Imputer
Batch size
256
256
256
[64, 128, 256)
Learning rate
le-3
le-4
le-4
[1e-3, le-4) [1, 1e3, le6) (100, 300, 500)
Regularization @ # Iterations
le3
le6
le6
500
500
500
Classifier
Hidden size (3-layer)
64
256
256
[64, 128, 256) [64, 128, 256) [le-3, 5e-4, le-4, 5e-5, 1e-5) (1,3,5,8, 10,12)
Batch size
256
256
256
Learning rate
5e-4
lc-4
le-4
Class weight
PPO
Learning rate
le-4
lc-4
le-4
(1e-3, le-4) [128, 256] (64, 128, 256) {1024, 2048)
Hidden size (2-layer)
128
256
128
Batch size
128
128
256
# Timesteps per update
1024
2048
2048
End-to-end
# Outer loop
100
100
[10,100)
# Epochs Classifier trained per loop
8
6
{2,4,6,8)
# Timesteps PPO trained per loop
5k
50k
50k
[5k, 50k
150
0.80
train
validation
Score
AUROC
Testing Cost
100
best
0.2
50
0.65
50
100
100
100
End-to-end training loops
End-to-end training loops
End-to-end training loops
Figure 5: RL training curves on AKI Dataset. We evaluate the policy and classifier learned during the end-to-end training phase both train, validation and test set. The matching curve confirms generalizability of the learned policies.
Ferritin Dataset
AKI Dataset
Sepsis Dataset
0.6
score
0.4
scare
0.4
0.4
0.2
0.2
0.0
0.0
0.0
25
50
100
50
100
150
50
100
150
Testing Cost
Testing Cost
Testing Cost
0.8
0.9
0.85
0.7
AUROC
AUROC
AUROC
0.80
0.8
0.6
0.5
0.75
Tacta
Dar
Do (K2ymell
0.7
200
0.4
0.70
25
50
75
100
50
100
150
0
50
100
150
Testing
Testing
Figure 6: Cost-F, Pareto Front for maximizing F1-score Ferritin, AKI and Sepsis Datasets
18

---

Published as a conference paper at ICLR 2023
C
EMFLOW
In this section, we describe the EMFlow imputation algorithm, originally proposed in Ma & Ghosh (202T), used in constructing the posterior state encoder. The key idea of EMFlow is to first map the data onto the latent space via a normalizing flow model, such that it enjoys Gaussian distribution in the latent space. Then it apply an online version of the expectation maximization algorithm in the latent space. The high-level picture of EMFlow is illustrated in Fig. 7 In each training iteration, we
X in
Initial Imputation
Data Space X
X
NF f
NF (.)
Z
Online EM
Data Space Z
Figure 7: Illustration of EMFlow Ma & Ghosh 2021
first fix the estimate base distribution (áù Σ) in the latent space and update the normalizing flow (NF) via optimizing the negative log-likelihooe of a batch B of the current imputation result:
(0)
log (&;; 0, i, E.
B
iEB
where NF is parametrized by 9 and denotes the likelihood of the current imputation estimate X; given the base distribution. Then with the current estimate of the normalizing flow, we do an online version of expectation maximization (EM) in the latent space to update the estimate for the base distribution. The expectation step given as:
m E 27: p, E1, i e B
(5)
And the maximization step follows iterative update scheme to deal with the memory overhead of the classical EM:
(1) PaPMhatch + (1 (--2
(6)
(00) Ebotch + (1 - pt) (t--1)
(7)
where Pbatch and Batch are estimated using the conditional mean and variance given a batch of data:
(Aattch gu 50); (z, MrheB)
(8)
Etntch = ()) 5011; {zi, MrJen)
Lastly, re-optimize NF via combination of log- likelihood of current imputation and regulariza- tion term indicating the distance of the imputation result to the ground truth.
(0
Σ [logpx (K;; p, E) a x
(10)
B
iEB
In this work, we used the above "supervised" version of EMFlow, as we already know the ground truth data when training this imputer. Finally, the imputation procedure is in similar manner as
19

---

Published as a conference paper at ICLR 2023
Algorithm 2 Training Phase of EMFlow ("supervised" version)
Input: Current imputation X = (8), X2, xx)) T. missing masks M , ",,"
initial estimates of the base distribution (á(00, (0))), online EM step size sequence: Pl, P
for t 1,2, Tepach do
Get mini-batch X = {$:
# update the flow model Compute L1 in Eq. (4)
Update via gradient descent # update the base distribution
Z, fo (x; E
Impute in the latent space with (a(t-1), ((--)) (2(--1),{((-1)) to get {2iJieB via Eq. (5)
Obtain updated (ji(0) _)))) via Eq. (6
# update the flow model again
X, fe (2.),
B
Compute Le via Eq. 10
Update 0 via gradient descent
end for
Algorithm Re-immptation Phase of EMFlow
Input: Current imputation X ( x2, xn, T. missing masks M (M1, Me, M,
initial estimates of the base distribution (Ø, Σn from the previous training phase for ( = 1,2, ni do
2, f (x)
Impute in the latent space with (f(t-11,, E(1-1)) to get {2i via Eq. (
x; -2, B
8; x; ① +
end for
Output: {
20

---

Published as a conference paper at ICLR 2023
the training phase without updating the parameters. Both the training and imputation scheme
illustrated in Algorithm and respectively.
To illustrate the imputation quality of the above EMFlow algorithm, present simple results in Table of the root mean square error (RMSE) compared to classical imputation algorithms on UCI datasets as reference. We do point out unlike conventional usage of imputation, use EMFlow mainly as posterior state encoder, instead of a imputer with excellent imputation quality.
Table 8: Comparison of Imputers on Public Datasets in RMSE. Each dataset has 20% MCAR. Apart from the time-consuming MissForest Imputer, the EMFlow outperforms other imputers in most datasets.
Imputer
Letter
Spam
News
Air
Credit
Mean
0.1535
0.0573
0.2265
0.1877
0.1592
MICE
0.1148
0,0664
0.1500
0.1014
0.1862
MissForest
0,0608
0.0476
0.1365
0.0869
0.1430
Matrix Completion
0.1290
0.0528
0.1653
0.1261
0.1632
GAIN
0.1198
0.0513
0.1441
0.1318
0.1858
EMFlow
0.1111
0.0776
0.1395
0.1022
0.1410
D END-TO-END TRAINING
In this section, we give more details of the proposed end-to-end training approach discussed in We give out specific form for the objective training both the classifier module and the panel/prediction selector module.
Training objective & for the classifier module The classifier fo(-) RR RC, parameterized by o, maps the posterior encoded state Impa .)) to a probability distribution over labels. It is trained by directly minimizing the cross entropy loss le defined as:
2
exp (I M)); Σ exp fo (Imp, M));
lo x, M) = Σ wi log
from the collected dataset Q in training the RL policy. This differs from typical classification in that the data are collected adaptively by RL rather than sampled from a prefixed source. Here w, 's are the class weights which treated as a tuning parameter. In this work, we used 3-layer DNN as the model for the classifier.
Training objective e,t for the panel/prediction selector module The panel selector, network module parameterized by sb, takes the following as input
cmb
$0.0 C M) (Impa @ M), (Imp, M)).
(11)
and maps it to a probability distribution over actions. We train the panel selector by using the classical proximal policy updates (PPO) Schulman al. (017)) which aims at maximizing clipped surrogate objective regularized by square-error loss of the value functions and an entropy bonus. We denote this loss function as lrt, which is defined the following regularized clipped surrogate loss Schulman et al. (017);
t gold) (CLIP(bi gold) cLVv(()) cg
where fCLiP(W; f/ gold) R. min( To (ail))
"gate (a. A, clip( Ne (a (8:) e, e) denotes the clipped surrogate loss with A, being the estimated advantages, regularized by the square error loss of the value function (v²) = E, (VV(
(o.
V and an entropy term Ent[TT.] over states. Here E, is empirical average over collected dataset Q, and & eemb (st) denotes the state embedding derived. Note the policy and value network share the parameter tb. We also show high-level algorithm There. For more details, please refer to the original paper (Schulman et al. (2017). We use the Python package 3table-baseline3 Raffin et al. (2021) for implementing PPO.
21

---

Published as a conference paper at ICLR 2023
Algorithm 4 Proximal Policy Optization (PPO)
for iteration = 0, 1, ... do
for actor 1, 2, .... , do
Run policy "you in environment for 7 timesteps and save all observations in Q P Q is also used to train the classifier module
Compute advantage estimate A.,... AA end for Optimizae surrogate t, wrt uh with K epochs and minibatch size M s Nactor Vold V
end for
E PROOFS OF THEOREM 4.1
For the ease of reading, we present the T'eorrm 4.1 in Appendix again.
Theorem E.1 (Copy of Theorem 4.1)
The Cost-Fi Pareto front defined in ① is a subset of the collection of all reward-shaped solutions, given by
II C II := Ux20,e50 argmax (7N(T) X TP(x) + p- Cos((TT).
As natural corollary, we have the following result:
Corollary E.1. The Cost-F; F2 Pareto front defined in is a subset of the solutions of the MDP model for dynamic diagnosis process defined in Section with reward functions:
p c(a),
if a E [D] (choosing task panels)
R((s,a) = A. 1y P),
if = P (true positive diagnosis)
1 if N (true negative diagnosis)
for all pairs of (X, p). Here y denotes the ground-truth label of the patient.
Before presenting the proofs for Theorem 4.1
reintroduce some key notations. First, define the length of the MDP episode as 7 min 2 Ola, E (P, N}} as random variable depending on the policy æ. Then
rigorously define the normalized true positive, negative, false positive
and false negative as
TP(x) = g P). TN(T) = 9N, N), FP(T) == N),,N), FN(T) == g(N, P).
Here ( == (a, = i} 1(y = j denotes the probability of the policy diagnosing class
i (P,N) while the ground truth label is of class j (P, N). And the testing cost is defined
Cost(T)= EF
Σ
c(k) 1aa = 1:)
RE(D]
where e(k) the cost of panel
The cumulative state-action occupancy measure μ : - RSXA 20
is defined as
1 (s, :- E Σ 1=0
Vs S
which denotes the expected time spent in state-action pair (s, a) during episode. There is one-one correspondence between the policy and the occupancy measure given by
p(s,a) a)
क (a|s)
Eeca ())
Proof of Theorem 4.7 We follow the proof framework stated in the main text.
22

---

Published as a conference paper at ICLR 2023
Step 1: utilizing monotonicity of Fi score To start with, note that Fi score is monotonically increasing in both true positive and true negative. Assume for any given cost budget B, the optimal policy x²
x* (B) achieves the highest Fj score. Then (B) is also the optimal to the following program:
max (T) subject to Cost(a) (T) s B, (P(T) 2 TP(x
indicating the Pareto front of P1 score is a subset of
II" C CUBS0,B'Ejo.1 argmaax (T)(() subject to Cost(T) TP(m() 2 B'..
(12)
Step 2: reformulation using occupancy measures Fix any specific pair (B, B'). Consider the equivalent dual linear program form Zhang et al. (2020) of the above policy optimization problem in terms of occupancy measure. Then the program (12 is equivalent
max TN(u)
subject to Cost(u) s B, TP(μ) 2 B', Σ ((s s'aa
Σ u d')P(ssssss ,') + Ess),
where E(-) denotes initial distribution, and TP, TN and cost are reloaded in terms of occupancy as:
TP(p) = E
p(s,a), a), TN(A)= =
Σ s.y=N,a=N
u(s, Cost(μ)= = Σ
c(k)) Σ u(s,a).
s.p=P.=PP
KEID)
s,amk
Step 3: utilizing hidden minimax duality a max-min program:
The above program be equivalently reformulated as
max min A20,p50
TN(1) + (TP(p) B') p. (Cost( - B)
subject to [m(s,a)= Σ pls', a(s'.a')PP d'ss(s a') + Ess).
s'a'E[D]
Note the max-min objective is linear terms of X, and μ. Thus minimax duality holds, then we swap the min and max to obtain the equivalent form:
min max TN(H) + A (TP(H) B") + p. (Cost(u) - B) x20,p50
#
subject Σm(s,o)= Σ pls' a P d') Es.).
d',a'E[D]
For any fixed pair of (A, p). the inner maximization problem of the above can be rewritten back in terms of policy π,, as the following equivalent unconstrained policy optimization problem:
max TN(T) ) TP(*) + Cost((r).
This is finally a standard cumulative-sum MDP problem, with reshaped reward: reward p c(t) for the action of choosing test panel f, reward À for the action to diagnose and get a true positive, reward 1 for getting a true negative, as stated in Corollary
(()),)),
if a E [D] (choosing task panels)
R((s.a) = 1 P}.
if a P (true positive diagnosis)
11y -N,,
if a N (true negative diagnosis)
23

---

Published as a conference paper at ICLR 2023
F EXTENSION TO AM METRIC
In this section, we show that our framework directly handles classic linear metrics defined as a linear combination (normalized) true positives and true negatives. We use the AM metric Natarajan et al. (2018): Menon et al. 2013), linear metric for imbalanced tasks via considering the average of true positive rate and true negative rate, an example.
We can formally define the cost-AM Pareto front as follows:
Am = UBSD argmax (AM(T)) subject to Cost(T) S B}
where AM(=) (TPR(T) + (T). Here we define the true positive rate and true negative rate respectively as:
TP(x)
TN(x)
TPR(*) TP(x) FN(T)'
TNR(")
TN( + FP(T)'
where (P(T), TN(T), FT()), (T normalized true positive, true negative, false positive and false negative rates that sum up to Let X > 0 be the known ratio between the number of healthy patients and ill patients. Then we have TP (T) (T) 1+x and + A_ indicating
AM(T)
(T)((T) + TN(T))
2A
Thus, we show the linearity of AM metric, indicating given any fixed budget, the problem already reduced back to standard MDP. It is simpler problem and directly solve them using training framework. Rigorously, we have the following result parallel to our Theorem 1:
AM = Upoo argmax (TN((T) TP(x) + Cost(T))
Empirically, show the following Table of results AM metric, comparing the same baselines discussed in our work. The testing costs are similar to the Table 2 in the paper. It is shown that our approach achieves the best AM score while having lower testing costs.
Table 9: Model performance measured by balanced accuracy for three real-world clinical dataset.
The model yielded the highest balanced accuracy in bold.
Models
AKT
Ferritin
Sepsis AM
Metrics
AM
Cost
AM
Cost
Cost
LR RF
0.714
$591
0.842
$290
0.752
$591
0.668
$591
0.817
$290
0.722
$59
XGBoost
0.644
$591
0.763
$290
0.644
$591
LightGBM
0.690
$591
0.832
$290
0.751
$591
LR (2 panels)
0.714
$92
0.783
$92
0.746
$92
RF (2 panels)
0.660
$92
0.778
0.707
$92
XGBoost (2 panels)
0.639
$92
0.683
$92
0.642
$92
LightGBM panels)
0.688
$92
0.786
$92
0.737
$92
SM-DDPO
0.741
$92
0.845
0.754
$48
G PRECISION AND RECALL
We also show the detailed statistics two sets of optimal solutions that optimizes Fl score and AM metric respectively. It can be shown by comparing Table 10 and that F1 score balances the recall and precision while AM metric balances the true positive and negative rates.
24

---

Published as a conference paper at ICLR 2023
Table 10: Detailed statisties of our results optimizing score.
Dataset
Recall
Precision
F1
true positive
true negative
AM
Ferritin
0.664
0.587
0.624
0.664
0.954
0.809
AKI
0.537
0.459
0.495
0.537
0.875
0.706
Sepsis
0.573
0.551
0.562
0.573
0.921
0.747
Table 11: Detailed statistics of results optimizing AM metric.
Dataset Ferritin
Recall
Precision
Fi
true positive
true negative rate
AM
0.798
0.420
0.550
0.798
0.892
0.845
AKI 0.704
0.385
0.498
0.704
0.778
0.741
Sepsis
0.613
0.497
0.549
0.613
0.895
0.754
25


Reviews:
[{"review_id": "nfZEDXLKliX", "reviewer": "Reviewer_Yh4E", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper investigates the use of reinforcement learning for lab test panel optimization, aiming to dynamically prescribe test panels based on available observations, to maximize diagnosis/prediction accuracy while keeping testing at a low cost. Given that clinical diagnostic data are often highly imbalanced, the authors maximize the F1 score instead of the error rate, and further develop a reward-shaping approach by leveraging the duality of policy optimization, so that the problem can be solved by standard RL methods. ", "summary_of_the_review": "The presentation is clear. The proposed method is novel and sound to me. Although the empirical performance is better than other methods in most cases, the F1 score still seems pretty low, which may still not be appropriate for clinical use.\n", "strength_and_weaknesses": "Strength:\n1. The studied problem is very interesting and practically useful, and I think using reinforcement learning for lab test panel optimization is an appropriate way.\n2. The proposed method seems sound to me (but I did not check the proofs).\n3. The writing is good.\n\nWeakness:\n1. The reported empirical performance of the proposed method has an improvement in most cases, compared to other methods. However, the F1 score still seems pretty low, which may still not be appropriate for clinical use.\n\n2. I would suggest the authors pay attention to the format, avoiding using too many \\vspace{}; some parts can be moved to the appendix.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "M4Xx4C41GWn", "reviewer": "Reviewer_Ws9d", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies sequential selection of laboratory test panels, formulated as a cost-sensitive adaptive feature selection/acquisition problem. The two main innovations are: \n1. For diagnostic performance, instead of maximizing accuracy, this work proposes to maximize F1 score to address high class imbalance (how high ). Since the F1 score cannot be written as an average of episode returns and thus cannot be directly formulated as a reward function, the paper presents new theoretical insights on transformations of the optimization problem which led to a reward shaping approach that allows RL to be applied. \n2. Viewing the problem as multi-objective optimization of diagnostic performance and cost, the proposed reward function is able to identify a superset of all Pareto optimal solutions, allowing one to recover the Pareto front of optimal F1 for each cost.\n\nExperiments were conducted on real-world datasets - ferritin, sepsis and AKI - and achieves good diagnostic performance while having lower costs than competitors. ", "summary_of_the_review": "The problem studied is interesting and relevant, the proposed solution makes sense. Certain claims need refinement and experiments need some clarifications. ", "strength_and_weaknesses": "**Strengths:**\n- The problem studied by this paper is interesting and important for many clinical classification problems subject to various budget constraints. \n- Novel theoretical insights on the how the optimization problem can be transformed to formulate a reward function that is an average of episode rewards.\n- Empirical results are strong, the proposed approach shows favorable performance on three different datasets. \n\n**Weaknesses:**\n- The title is missing keywords such as \"cost-sensitive\" and \"active feature acquisition\" (\"diagnostic policy learning\" might be a new word to many readers). \n- In Fig 5, some of the red/blue points lie outside the Pareto front. How are the Pareto front curves fitted? \n- Some statements could use a bit more clarification:\n  - in Step 3, $(s',a')$ with the apostrophe is typically used to indicate next state-action pair, where here in the occupancy definition they represent the previous state-action pair. It is suggested that the authors change the notation here to something else like $(\\tilde{s}, \\tilde{a})$ throughout. \n  - in the abstract, \"the F1 score cannot be formulated as a simple sum of cumulative rewards\" - on first read this sentence is confusing, it seems to be discussing the cumulative rewards of a single trajectory, but the F1-score is ill defined at the trajectory/instance level. Only after reading the paper did I understand it's actually talking about the sum (or average) of multiple trajectories. I suggest this sentence be rephrased. \n  - on page 4, Step 2: \"Fix any specific pair (B,B')\" should be \"Fix any specific pair (B,K)\"\n  - in Sec 4 the definition of TP(\\pi) vs TP(\\mu) seems inconsistent on their normalization. Are they the number of true positives, or the true positive rate? \n  - in Sec 3.1 and Fig 1 \"A penalty will [be] generated if the diagnosis does not match the ground truth $y$.\" This sentence is inconsistent with Sec 4 Step 3 reward definition $\\mathrm{TN} + \\lambda \\mathrm{TP} + \\rho \\mathrm{Cost}$ where no penalty is assigned, only positive rewards for true pos/neg. \n  - On page 5: regarding \"semi-model-based\", could you clarify why the reward needs to be learned? It seems that the reward r(s,a) should only depend on the diagnosis prediction and the true label. \n  - Sec 3.1 should also clarify that the transition dynamics of the MDP can be perfectly specified. \n   - in Sec 5.3: Eqn (3) is not only the panel selector, it's panel/prediction selector which is the policy network. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "8TjztwvBZ4", "reviewer": "Reviewer_i4nR", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposed a reinforcement learning-based approach for test panel selection in order to make diagnostic predictions under cost constraints. The paper deals with data imbalance in the data by a rewarding shaping approach. It proposed a semi-model-based approach to learning the diagnostic policy. Experiments are carried out on three real-world datasets to demonstrate the utility of the proposed method.", "summary_of_the_review": "Overall, I am leaning toward accepting the paper because the paper deals with an interesting problem in ML for health, formulates and develops a reasonable solution to the problem, and evaluates the proposed method rigorously across multiple real-world datasets over multiple baseline methods.\n\nMy concerns are centered around the necessity of using the F1 metric, the interpretability of the proposed method given that the methods address problems in the healthcare domain, and the multi-objective formulation of the problem. It may also be desirable for the authors to discuss the applicability of alternative metrics for imbalanced classification such as the AM metrics and potentially carry out empirical comparison.", "strength_and_weaknesses": "Strength:\n1. the paper deals with an important problem in machine learning for health care.\n2. The paper proposes a semi-model-based RL approach to learning the diagnostic policy, which seems suitable to the problem formulation.\n3. Experiments are carried out on three real-world datasets, which is a fair amount of different problem settings. The proposed method is compared to 12 alternatives, which is fairly exhaustive.\n\nWeakness:\n1. It is not clear to me how necessary it is to directly optimize for F1 via reward shaping, given that there are alternative linear metrics available for imbalanced classification such as the AM metric. See, for examples,\nhttps://jmlr.org/papers/v18/15-226.html\nhttps://proceedings.mlr.press/v28/menon13a.html\n\n2. It is not clear to me how interpretable the diagnostic policy might be. This is in contrast to the tree-based baseline methods compared to the proposed methods.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster