Original: UnfoldML :A Cost-Aware 2-D Dynamic  PredictionPipeline for Multi-Stage Classiﬁcation Anonymous Author(s)
Revised: UnfoldML : Cost-Aware and Uncertainty-Based Dynamic 2D Prediction for Multi-Stage Classification

Original: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: Machine Learning (ML) research has focused on maximizing the accuracy ofpredictive tasks.
Revised: Machine Learning (ML) research has focused on maximizing the accuracy of predictive tasks.

Original: ML models, however, are increasingly more complex, resourceintensive, and costlier to deploy in resource constrained environments.
Revised: ML models, however, are increasingly more complex, resource intensive, and costlier to deploy in resource-constrained environments.

Original: These issuesare exacerbated for prediction tasks with sequential classiﬁcation on progressivelytransitioned stages with “happens-before” relation between them.
Revised: These issues are exacerbated for prediction tasks with sequential classification on progressively transitioned stages with “happens-before” relation between them.

Original: We argue that itis possible to “unfold” a monolithic single multi-class classiﬁer, typically trainedfor all stages using all data, into a series of single-stage classiﬁers.
Revised: We argue that it is possible to “unfold” a monolithic single multi-class classifier, typically trained for all stages using all data, into a series of single-stage classifiers.

Original: Each singlestage classiﬁer can be cascaded gradually from cheaper to more expensive binaryclassiﬁers that are trained using only the necessary data modalities or featuresrequired for that stage.
Revised: Each singlestage classifier can be cascaded gradually from cheaper to more expensive binary classifiers that are trained using only the necessary data modalities or features required for that stage.

Original: UnfoldML is a cost-aware 2-D dynamic  prediction pipelinefor multi-stage classiﬁcation that achieves this enabling (1) navigation of theaccuracy/cost tradeoff space, (2) reducing the spatio-temporal cost of inferenceby orders of magnitude, and (3) early prediction on proceeding stages.
Revised: UnfoldML is a cost-aware and uncertainty-based dynamic 2D prediction pipeline for multi-stage classification that enables (1) navigation of the accuracy/cost tradeoff space, (2) reducing the spatio-temporal cost of inference by orders of magnitude, and (3) early prediction on proceeding stages.

Original: UnfoldMLachieves orders of magnitude better cost in clinical settings, while detecting multistage disease development in real time.
Revised: UnfoldML achieves orders of magnitude better cost in clinical settings, while detecting multistage disease development in real time.

Original: It achieves within 0.1% accuracy fromthe highest-performing multi-class baseline, while saving close to 20X on spatiotemporal cost of inference and earlier (3.5hrs)
Revised: It achieves within 0.1% accuracy from the highest-performing multi-class baseline, while saving close to 20X on spatiotemporal cost of inference and earlier (3.5hrs)

Original: disease onset prediction.
Revised: disease onset prediction.

Original: We alsoshow that UnfoldML generalizes to image classiﬁcation, where it can predictdifferent level of labels (from coarse to ﬁne) given different level of abstractionsof a image, saving close toa 5X cost with as little as 0.4% accuracy reduction.
Revised: We also show that UnfoldML generalizes to image classification, where it can predict different level of labels (from coarse to fine) given different level of abstractions of a image, saving close to 5X cost with as little as 0.4% accuracy reduction.

Original: UnfoldML improves the whole frontier of optimality in the accuracy/cost tradeoffspace, achieving the highest area under the curve.

Original:  Introduction
Revised:  Introduction

Original: Machine Learning (ML) research has mostly focused on improving prediction accuracy for classiﬁcation tasks, such as image classiﬁcation (Foret et al., 2020; Xie et al., 2017), disease risk prediction(Feng et al., 2020; Xu et al., 2018),
Revised: Machine Learning (ML) research has mostly focused on improving prediction accuracy for classification tasks, such as image classification (Foret et al., 2020; Xie et al., 2017), disease risk prediction

Original:  pedestrian detection Zhang et al., 2018; Cai et al., 2015), etc.
Revised: (Feng et al., 2020; Xu et al., 2018), pedestrian detection (Zhang et al., 2018; Cai et al., 2015), etc.

Original: Theunderstandable drive for high accuracy has often resulted in deeper, complex neural networks, which can incur high memory ( spatio cost) and high latency ( temporal cost) at inference time.
Revised: The understandable drive for high accuracy has often resulted in deeper, complex neural networks, which can incur high memory ( spatial cost) and high latency ( temporal cost) at inference time.

Original: However, the deployment of ML applications must be cost-aware.
Revised: However, the deployment of ML applications must be cost-aware.

Original: Run time environment like mobile devices or bedside-patient monitors are commonly resource constrained, and applications that can be ofﬂoaded to cloud computing always aim for a reduced cloud bill.
Revised: Run time environment like mobile devices or bedside-patient monitors are commonly resource constrained, and applications that can be offloaded to cloud computing always aim for a reduced cloud bill.

Original: In this paper, we focus on developing a pipeline that can balance between high prediction accuracy and low spatio-temporal cost in deploying
Revised: In this paper, we focus on developing a pipeline that can balance between high prediction accuracy and low spatio-temporal cost in deploying

Original: ML classiﬁcation models
Revised: ML classification models.

Original: We consider a scenario of deploying ML classiﬁers in a multi-stage classiﬁcation task where one predicted class can progressively transform to a next stage of classes, characterized as “happens- before” relationship between classes.
Revised: We consider a scenario of deploying ML classifiers in a multi-stage classification task where one predicted class can progressively transform to a next stage of classes, characterized as “happensbefore” relationship between classes.

Original: This task is commonly observed in many real-world applications.
Revised: This task is commonly observed in many real-world applications.

Original: For instance in clinical settings, disease progression is often identiﬁed by a series of stage transitions
Revised: For instance in clinical settings, disease progression is often identified by a series of stage transitions

Original: Abstract Input
Revised: Abstract Input

Original: Detailed InputIDK IDK ICK
Revised: Detailed Input

Original: [Figure]
Revised: [Figure]

Original: Exit

Original: [Figure] C a s [Figure] d e s u pg r ad e t
Revised: C a s ca d e s u pg r ad e t

Original: o c o s t l i e r m od e l s w h e n l ac 
Revised: o c o s tli e r m od e l s w h e n l ac k

Original: k i ng c on f i d e n ce
Revised: i ng c on f i d e n ce

Original: (I D K )
Revised: (I D K )

Original: [Figure]
Revised: [Figure]

Original: ICK
Revised: StageStage- Stage

Original: [Figure]
Revised: [Figure]

Original: ICK

Original: [Figure]
Revised: [Figure]

Original: ICK

Original: ICK [Figure] ICK ICK
Revised: (Sperling et al., 2011; Singer et al., 2016).

Original: ICK [Figure] [Figure] appropriate actions in time before it enters into late severe stages.
Revised: Detection on early stage of the disease allows doctors to take appropriate actions in time before it enters into late severe stages.

Original: In image classiﬁcations, recognitionfrom super classes to sub classes in a coarse-to-ﬁne manner has shown improved classiﬁcationperformance (Dutt et al., 2017; Lei et al., 2017).
Revised: In image classifications, recognition from super classes to sub classes in a coarse-to-fine manner has shown improved classification performance (Dutt et al., 2017; Lei et al., 2017).

Original: In all these applications, general multi-classclassiﬁers (Liu et al., 2019; Fagerström et al., 2019; Foret et al., 2020) have been developed bytreating all the stages as multi classes and achieved state-of-the-art prediction performance, but atsigniﬁcantly high spatio-temporal cost.
Revised: In all these applications, general multi-class classifiers (Liu et al., 2019; Fagerström et al., 2019; Foret et al., 2020) have been developed by treating all the stages as multi classes and achieved state-of-the-art prediction performance, but at significantly high spatio-temporal cost.

Original: Prior work (Wang et al., 2017) trades off accuracy for lowcost but still ignores the key relationship between the classes, so it fails to ﬁnd the optimal trade-off.
Revised: Prior work (Wang et al., 2017) trades off accuracy for low cost but still ignores the key relationship between the classes, so it fails to find the optimal trade-off.

Original: We propose UnfoldML  : a cost-aware  prediction pipeline that leverages the “happens-before” propertyin multi-stage classiﬁcation.
Revised: We propose UnfoldML 2 : a cost-aware and uncertainty-based prediction pipeline for dynamic multistage classification.

Original: It “unfolds” a monolithic multi-class classiﬁer into a series of singlestage classiﬁers, reducing its deployment cost.
Revised: It “unfolds” a monolithic multi-class classifier into a series of single-stage classifiers, reducing its deployment cost.

Original: Each single-stage classiﬁer is then cascaded graduallyfrom cheaper to more expensive binary classiﬁers, further reducing the cost by dynamically selectingan appropriate classiﬁer for an input query.
Revised: Each single-stage classifier is then cascaded gradually from cheaper to more expensive binary classifiers, further reducing the cost by dynamically selecting an appropriate classifier for an input query.

Original: Figure 1 summarizes the two dimensional (2-D) querypropagation mechanism designed in UnfoldML : Horizontally it allows a query to transition throughmultiple stages, and vertically it allows the query to progressively upgrade to costlier modelsconstrained by the pre-speciﬁed budget limit.
Revised: Figure 1 summarizes the two dimensional (2D) query propagation mechanism designed in UnfoldML : Horizontally it allows a query to transition through multiple stages, and vertically it allows the query to progressively upgrade to costlier models constrained by the pre-specified budget limit.

Original: It computes a classiﬁer’s prediction conﬁdenceon a query then directs the query through one the following three gates: 1) “ I conﬁdently know
Revised: It computes a classifier’s prediction confidence on a query then directs the query through one the following three gates: 1) “ I confidently know

Original: NO ” ( ICK 0 ), which rejects the current query and early exits from the pipeline (exit); 2) “ I don’tknow ” ( IDK ), which upgrades the query to a higher accuracy but costlier model, producing a moreconﬁdent result within the budget (vertical cascading); and 3) “ I conﬁdently know YES ” ( ICK 1 ),which transitions the query to the next stage of the prediction task (horizontal forwarding).
Revised: NO ” ( ICK 0 ), which rejects the current query and early exits from the pipeline (exit); 2) “ I don’t know ” ( IDK ), which upgrades the query to a higher accuracy but costlier model, producing a more confident result within the budget (vertical cascading); and 3) “ I confidently know YES ” ( ICK 1 ), which transitions the query to the next stage of the prediction task (horizontal forwarding).

Original: The designof ICK 0 gate allows queries to early exit from the UnfoldML so it reduces the overall spatio-temporalcost of the pipeline.
Revised: The design of ICK 0 gate allows queries to early exit from the UnfoldML so it reduces the overall spatio-temporal cost of the pipeline.

Original: The ICK 1 gate allows queries to faster transition to next stages so it enablesearly prediction on late stages, which can be critical in clinical settings (Reyna et al., 2019).
Revised: The ICK 1 gate allows queries to faster transition to next stages so it enables early prediction on late stages, which can be critical in clinical settings (Reyna et al., 2019).

Original: Overall,the combined 2-D propagation mechanism uniquely enables the navigation of the cost/accuracytradeoff space for searching an optimal set of policies for dynamic model selection at inference time.
Revised: Overall, the combined 2D propagation mechanism uniquely enables the navigation of the cost/accuracy tradeoff space for searching an optimal set of policies for dynamic model selection at inference time.

Original: We also propose two training algorithms for learning the optimal policies for the designed 2-Dquery propagation in UnfoldML .
Revised: We also propose two training algorithms for learning the optimal policies for the designed 2D query propagation in UnfoldML .

Original: The key idea of the training algorithms is to learn the optimalthresholds on the gating functions deﬁned for ICK 0 , IDK and ICK 1 .
Revised: The key idea of the training algorithms is to learn the optimal thresholds on the gating functions defined for ICK 0 , IDK and ICK 1 .

Original: The ﬁrst proposed hardgating algorithm assumes the gating functions to be step functions parameterized by deterministicconﬁdence thresholds.
Revised: The first proposed hardgating algorithm assumes the gating functions to be step functions parameterized by deterministic confidence thresholds.

Original: To ﬁnd the optimal thresholds, it performs bottom-up grid search over atopologically-sorted list of all the models and identiﬁes the thresholds to minimize the prediction loss,while following a cost constraint.
Revised: To find the optimal thresholds, it performs bottom-up grid search over a topologically-sorted list of all the models and identifies the thresholds to minimize the prediction loss, while following a cost constraint.

Original: The secondly proposed soft-gating algorithm deﬁnes the gatingfunctions to be probabilistic activation functions.
Revised: The secondly proposed soft-gating algorithm defines the gating functions to be probabilistic activation functions.

Original: It follows a Mixture-of-Expert (MoE) frameworkto adaptively determine a models’ conﬁdence threshold given all the other model’s conﬁdence inpredicting the same query.
Revised: It follows a Mixture-of-Expert (MoE) framework to adaptively determine a models’ confidence threshold given all the other model’s confidence in predicting the same query.

Original: To obviate the cost of running all models in MoE selection, we furtherpropose a Dirichlet Knowledge Distillation (DKD) to run only a cheap multi-label classiﬁer that istrained for distilling the Bayesian predictive uncertainties of all models.
Revised: To obviate the cost of running all models in MoE selection, we further propose a Dirichlet Knowledge Distillation (DKD) to run only a cheap multi-label classifier that is trained for distilling the Bayesian predictive uncertainties of all models.

Original: We summarize the contributions of this paper as follows:
Revised: We summarize the contributions of this paper as follows:

Original: • We design a novel 2-D query propagation pipeline that “unfolds” multi-stage prediction workﬂowsby leveraging the “happens-before” relationship between the stages, and achieves a lower-cost prediction pipeline with minimal accuracy degradation.
Revised: • We design a novel 2D query propagation pipeline that “unfolds” multi-stage prediction workflows by leveraging the “happens-before” relationship between the stages, and achieves a lower-cost prediction pipeline with minimal accuracy degradation.

Original: • We propose two learning algorithms to sufﬁciently navigate the cost/accuracy tradeoff space and search an optimal set of policies for the designed 2-D query propagation.•
Revised: • We propose two learning algorithms to sufficiently navigate the cost/accuracy tradeoff space and search an optimal set of policies for the designed 2D query propagation.

Original:  We apply the proposed pipeline to two real-world applications and demonstrate it reduces thespatio-temporal cost of inference by orders of magnitude.
Revised: • We apply the proposed pipeline to two real-world applications and demonstrate it reduces the spatio-temporal cost of inference by orders of magnitude.

Original:  Related Work
Revised:  Related Work

Original: The most relevant work to our proposed method is the one-step IDK cascade (Wang et al., 2017),which incorporates prior work of “ I don’t know ” ( IDK ) classes (Trappenberg and Back, 2000; Khaniet al., 2016) into cascade construction and introduce a latency-aware objective into the constructioncomparing with previous cascaded prediction frameworks (Rowley et al., 1998; Viola and Jones,2004; Angelova et al., 2015).
Revised: The most relevant work to our proposed method is the one-step IDK cascade (Wang et al., 2017), which incorporates prior work of “ I don’t know ” ( IDK ) classes (Trappenberg and Back, 2000; Khani et al., 2016) into cascade construction and introduce a latency-aware objective into the construction comparing with previous cascaded prediction frameworks (Rowley et al., 1998; Viola and Jones, 2004; Angelova et al., 2015).

Original: Another group of work focus on the problem of feature selectionassuming each feature can be acquired for a cost.
Revised: Another group of work focus on the problem of feature selection assuming each feature can be acquired for a cost.

Original: They train a cascade of classiﬁers for optimizingthe trade-off between the expected classiﬁcation error and the feature cost.
Revised: They train a cascade of classifiers for optimizing the trade-off between the expected classification error and the feature cost.

Original: Early solution (Raykaret al., 2010) limits the cascade to a family of linear discriminating functions.
Revised: Early solution (Raykar et al., 2010) limits the cascade to a family of linear discriminating functions.

Original: Cai et al.
Revised: Cai et al.

Original: (2015) appliesboosting method for cascading a set of weak learners.
Revised: (2015) applies boosting method for cascading a set of weak learners.

Original: Recent methods (Trapeznikov and Saligrama,2013; Clertant et al., 2019; Janisch et al., 2019) develop POMDP-based frameworks and incorporatedeep Q-learning in training the cascades.
Revised: Recent methods (Trapeznikov and Saligrama, 2013; Clertant et al., 2019; Janisch et al., 2019) develop POMDP-based frameworks and incorporate deep Q-learning in training the cascades.

Original: In contrast to all of the above work that are only 1-Dpipelines for one-step prediction task (can be multi-class classiﬁcations), our method extends to a2-D pipeline that can dynamically forward examples to next steps after they are conﬁdently predictedas passed on the current step.
Revised: In contrast to all of the above work that are only 1-D pipelines for one-step prediction task (can be multi-class classifications), our method extends to a 2D pipeline that can dynamically forward examples to next steps after they are confidently predicted as passed on the current step.

Original: Further, we also develop a more efﬁcient pipeline framework based on
Revised: Further, we also develop a more efficient pipeline framework based on

Original: Mixture-of-Experts (MoE) modeling and knowledge distillation, which can apply gradient decentalgorithms for learning the parameters efﬁciently.
Revised: Mixture-of-Experts (MoE) modeling and knowledge distillation, which can apply gradient decent algorithms for learning the parameters efficiently.

Original: The idea of MoE was originally introduced by Jacobs et al.
Revised: The idea of MoE was originally introduced by Jacobs et al.

Original: ([n. d.]), for partitioning the trainingdata and feeding them into separate neural networks during the learning process.
Revised: (1991), for partitioning the training data and feeding them into separate neural networks during the learning process.

Original: This gate decisiondesign is applied into many domains such as language modeling (Ma et al., 2018), video captioning(Wang et al., 2019), multi-tasking learning (Ma et al., 2018).
Revised: This gate decision design is applied into many domains such as language modeling (Ma et al., 2018), video captioning (Wang et al., 2019), multi-tasking learning (Ma et al., 2018).

Original: It is also used in network architecturesearching (Eigen et al., 2013) by setting gate activation on network layers.
Revised: It is also used in network architecture searching (Eigen et al., 2013) by setting gate activation on network layers.

Original: Sparse gates are introducedin MoE so that it can efﬁciently select from thousands of sub-networks (Shazeer et al., 2017) aswell as increases the representation power of large convolutional networks by only using a shallowembedding network to produce the mixture weights (Wang et al., 2020).
Revised: Sparse gates are introduced in MoE so that it can efficiently select from thousands of sub-networks (Shazeer et al., 2017) as well as increases the representation power of large convolutional networks by only using a shallow embedding network to produce the mixture weights (Wang et al., 2020).

Original: We incorporate the ideaof sparsely gated MoE (Shazeer et al., 2017; Wang et al., 2020) into our prediction framework, anddesign a soft-gating training algorithm by using ReLU as the sparse gating function and imposing
Revised: We incorporate the idea of sparsely gated MoE (Shazeer et al., 2017; Wang et al., 2020) into our prediction framework, and design a soft-gating training algorithm by using ReLU as the sparse gating function and imposing

Original: L1-norm regularization on the gating weights for further sparsity.
Revised: L1-norm regularization on the gating weights for further sparsity.

Original: Conﬁdence criterion has been incorporated into active learning by Li and Sethi (2006) and thenextended by Zhu et al.
Revised: Confidence criterion has been incorporated into active learning by Li and Sethi (2006) and then extended by Zhu et al.

Original: Lei (2014) proposed conﬁdence based classiﬁers that identiﬁes theconﬁdent region (like ICK class) and uncertain region (like IDK class) in predictions.
Revised: Lei (2014) proposed confidence based classifiers that identifies the confident region (like ICK class) and uncertain region (like IDK class) in predictions.

Original: Conﬁdenceare also introduced into word embedding (Vilnis and McCallum, 2015; Athiwaratkun and Wilson,2018) and graph representations (Orbach and Crammer, 2012; Vashishth et al., 2019).
Revised: Confidence are also introduced into word embedding (Vilnis and McCallum, 2015; Athiwaratkun and Wilson, 2018) and graph representations (Orbach and Crammer, 2012; Vashishth et al., 2019).

Original: Our methodposits thresholds on prediction conﬁdence for activating the gates in pipeline expansion.
Revised: Our method posits thresholds on prediction confidence for activating the gates in pipeline expansion.

Original: Bayesian
Revised: Bayesian

Original: Prior Networks (BPNs) (Malinin and Gales, 2018) have been proposed to estimate the uncertaintydistribution in model predictions, which is more computationally efﬁcient than traditional Bayesianapproaches (MacKay, 1992; Mackay, 1992; Hinton and Van Camp, 1993).
Revised: Prior Networks (BPNs) (Malinin and Gales, 2018) have been proposed to estimate the uncertainty distribution in model predictions, which is more computationally efficient than traditional Bayesian approaches (MacKay, 1992; Mackay, 1992; Hinton and Van Camp, 1993).

Original: We propose Dirichlet
Revised: We propose Dirichlet

Original: Knowledge Distillation (DKD) based on BPNs for distilling prediction uncertainty in large models sothat we only need to run a low-cost multi-head model for producing the weights in MoE efﬁciently.
Revised: Knowledge Distillation (DKD) based on BPNs for distilling prediction uncertainty in large models so that we only need to run a low-cost multi-head model for producing the weights in MoE efficiently.

Original:  Multi-Stage Dynamic Prediction Pipeline
Revised:  Multi-Stage Dynamic Prediction Pipeline

Original: We introduce a dynamic 2-D prediction pipeline UnfoldML , which learns the optimal policy formaking “ I conﬁdently know ” ( ICK ) predictions on sequential multi-stage classiﬁcation tasks.
Revised: We introduce a dynamic 2D prediction pipeline UnfoldML , which learns the optimal policy for making “ I confidently know ” ( ICK ) predictions on sequential multi-stage classification tasks.

Original: Anoptimal policy will effectively trade off prediction accuracy against spatio-temporal costs in order tomaximize the overall system accuracy’s AUC while staying under user-imposed cost constraints.
Revised: An optimal policy will effectively trade off prediction accuracy against spatio-temporal costs in order to maximize the overall system accuracy’s AUC while staying under user-imposed cost constraints.

Original: 3.1 Problem Formulation
Revised: 3.1 Problem Formulation

Original: Given x 2 X at time t , a multi-stage pipeline decides whether the individual should maintain atthe current stage s or progress into the next stage s + 1 for time t + 1 .
Revised: Given x ∈ X at time t , a multi-stage pipeline decides whether the individual should maintain at the current stage s or progress into the next stage s + 1 for time t + 1 .

Original: If there are a total numberof S stages that need to be detected, we train K number of models m for each speciﬁc stage s to form a model zoo M = {{ m 11 , · · · , m 1 K 1 } , · · · , { m S 1 , ..., m SK S }} .
Revised: If there are a total number of S stages that need to be detected, we train K number of models m for each specific stage s to form a model zoo M = {{ m 11 , · · · , m 1 K 1 } , · · · , { m S 1 , ..., m SK S }} .

Original: We measure each model’s132 spatio-temporal cost by multiplying the device cost per unit time with the serving time per predictionstage, denoted as cost ( m sk ) .
Revised: We measure each model’s spatio-temporal cost by multiplying the device cost per unit time with the serving time per prediction stage, denoted as cost ( m sk ) .

Original: To optimize the limited system resources, we design a 2-D UnfoldML in the following way: (1) startwith the simplest model for prediction of the initial stage on incoming data, and (2) upgrade verticallyto costlier models on those samples where “ I don’t know ” ( IDK ), or (3) transition horizontally tothe next stage of the pipeline samples where “ I conﬁdently know YES ” ( ICK 1 ) that we have correctlyidentiﬁed the sample at the current stage, otherwise (4) exit the pipeline for those “ I conﬁdently know
Revised: To optimize the limited system resources, we design a 2D UnfoldML in the following way: (1) start with the simplest model for prediction of the initial stage on incoming data, and (2) upgrade vertically to costlier models on those samples where “ I don’t know ” ( IDK ), or (3) transition horizontally to the next stage of the pipeline samples where “ I confidently know YES ” ( ICK 1 ) that we have correctly identified the sample at the current stage, otherwise (4) exit the pipeline for those “ I confidently know

Original: NO ” ( ICK 0 ) samples that have been identiﬁed and discarded.
Revised: NO ” ( ICK 0 ) samples that have been identified and discarded.

Original: Figure 1 (a) demonstrates the proposed2-D architecture of UnfoldML .
Revised: Figure 1 (a) demonstrates the proposed 2D architecture of UnfoldML .

Original: The central problem of UnfoldML is learning the optimal policy foreach of the three classes of gating functions with the objective of maximizing high system-wideaccuracy while minimizing the prediction cost.
Revised: The central problem of UnfoldML is learning the optimal policy for each of the three classes of gating functions with the objective of maximizing high system-wide accuracy while minimizing the prediction cost.

Original: We formulate UnfoldML as a decision rule mapping function m casc :
Revised: We formulate UnfoldML as a decision rule mapping function m casc :

Original: X ⇥ M ! M , which takesexample x t coming at time t and the current model choice m sk as an input to determine whetherthe model for the query should take one of the aforementioned three actions: upgrade vertically,transition horizontally, or exit the pipeline.
Revised: X × M → M , which takes example x t coming at time t and the current model choice m sk as an input to determine whether the model for the query should take one of the aforementioned three actions: upgrade vertically, transition horizontally, or exit the pipeline.

Original: These decisions rules can be realized by two groupsof parameters: a conﬁdence criterion q : X ⇥ M ! [0 , 1] that measures the conﬁdence score of amodel’s prediction on a data example, and two gate functions G IDK , G ICK 1 : [0 , 1] ⇥ { True , False }that are applied on to the conﬁdence score per each prediction made by model m sk .
Revised: These decisions rules can be realized by two groups of parameters: a confidence criterion q : X × M → [0 , 1] that measures the confidence score of a model’s prediction on a data example, and two gate functions G IDK , G ICK 1 : [0 , 1] × { True , False } that are applied on to the confidence score per each prediction made by model m sk .

Original: The two exclusivegates IDK and ICK 1 each respectively decides if the current prediction belongs to either an IDKclass such that (s.t.)
Revised: The two exclusive gates IDK and ICK 1 each respectively decides if the current prediction belongs to either an IDK class such that (s.t.)

Original: the system will upgrade the query to a costlier model while remaining withinthe user-deﬁned cost budget, or an ICK 1 class s.t.
Revised: the system will upgrade the query to a costlier model while remaining within the user-defined cost budget, or an ICK 1 class s.t.

Original: the system will transition the query to the nextstage in the pipeline.
Revised: the system will transition the query to the next stage in the pipeline.

Original: The third gate G ICK 0 is determined by ¬ G IDK ^ ¬ G ICK 1 .
Revised: The third gate G ICK 0 is determined by ¬ G IDK ∧ ¬ G ICK 1 .

Original: We formalize thedecision rule used at each prediction stage as follows
Revised: We formalize the decision rule used at each prediction stage as follows

Original: [Equation]
Revised: [Equation]

Original:  where q sk ( x t ) is a short notation for q ( x t , m sk ) measuring the conﬁdence of model m sk ’s predictionon data x t .
Revised: where q sk ( x t ) is a short notation for q ( x t , m sk ) measuring the confidence of model m sk ’s prediction on data x t .

Original: The goal of conﬁguring an optimal pipeline given a restricted computation resource canbe formalized as the following optimization problem:
Revised: The goal of configuring an optimal pipeline given a restricted computation resource can be formalized as the following optimization problem:

Original: [Equation]
Revised: [Equation]

Original: where G consists of the two gate functions G IDK and G ICK 1 , L denotes the end-to-end predictionloss on data D and c is a user-speciﬁed cost-constraint for the system.
Revised: where G consists of the two gate functions G IDK and G ICK 1 , L denotes the end-to-end prediction loss on data D and c is a user-specified cost-constraint for the system.

Original: 3.2 Gate Parameters Learning
Revised: 3.2 Gate Parameters Learning

Original: Given a training data set D , which does not include any data that was used in the training of the modelsin our model zoo: D = {  x i , ( y 1 i , t 1 i ) , · · · , ( y Si , t Si )  } Ni =1 , where x i =
Revised: Given a training data set D , which does not include any data that was used in the training of the models in our model zoo: D = { (cid:0) x i , ( y 1 i , t 1 i ) , · · · , ( y Si , t Si ) (cid:1) } Ni =1 , where x i =

Original: ( x i 1 , · · · , x iT i ) is an inputsequence observed for individual i , y si 2 { 0 , 1 } indicates whether the example entered to stage- s ,and if yes , we use t s i 2 ; [ [1 , T i ] to determine when it entered.
Revised: ( x i 1 , · · · , x iT i ) is an input sequence observed for individual i , y si ∈ { 0 , 1 } indicates whether the example entered to stage- s , and if yes , we use t si ∈ ∅ ∪ [1 , T i ] to determine when it entered.

Original: We ﬁrst partition the multi-stage datainto S one-stage data sets: D s = { ( x i [ t s  1 i : t si ] , y si == 1) } [ { ( x i [ t s  1 i : T i ] , y si == 0) } , then dividethe learning of IDK and ICK gate parameters into two separable sub-problems:
Revised: We first partition the multi-stage data into S one-stage data sets: D s = { ( x i [ t s − 1 i : t si ] , y si == 1) } ∪ { ( x i [ t s − 1 i : T i ] , y si == 0) } , then divide the learning of IDK and ICK gate parameters into two separable sub-problems:

Original: [Equation]
Revised: [Equation]

Original: where L s is the one-stage prediction loss on data D s , and c s is the cost budget that is pre-allocatedfor stage- s satisfying 
Revised: where L s is the one-stage prediction loss on data D s , and c s is the cost budget that is pre-allocated for stage- s satisfying (cid:80) s

Original: P s c s = c .
Revised: c s = c .

Original: Decomposing the end-to-end optimization problem in Eq.
Revised: Decomposing the end-to-end optimization problem in Eq.

Original: (1) intotwo sub-problems in Eq.
Revised: (1) into two sub-problems in Eq.

Original: (2) allows us to parallelize the training process.
Revised: (2) allows us to parallelize the training process.

Original: We can efﬁciently learneach stage’s optimal IDK gate parameters by solving Sub-Objective 1, and learn the optimal ICK 1gate parameters by ﬁxing the learnt IDK values in Sub-Objective 2.
Revised: We can efficiently learn each stage’s optimal IDK gate parameters by solving Sub-Objective 1, and learn the optimal ICK 1 gate parameters by fixing the learnt IDK values in Sub-Objective 2.

Original: 3.2.1 Hard-gating Training Algorithm
Revised: 3.2.1 Hard-gating Training Algorithm

Original: In this algorithm, we assume G IDK to be a hard-gating function that is parameterized by cutoff↵ sk s.t.
Revised: In this algorithm, we assume G IDK to be a hard-gating function that is parameterized by cutoff α sk s.t.

Original: the gate is only activated if the level of conﬁdence in the prediction is below a threshold.
Revised: the gate is only activated if the level of confidence in the prediction is below a threshold.

Original: Hard-gating is deﬁned as G IDK ( q sk ( x t ))

Original: =

Original: I ( q sk ( x t ) <

Original: ↵ sk ) ,

Original: [Equation]
Revised: [Equation]

Original: -gating

Original: [Equation]

Original: [Equation]

Original: where I ( · ) is an indicator function.
Revised: where I ( · ) is an indicator function.

Original: Based on the Problem Formulation, a model m sk at stage scan only be activated if I ( q sk ( x t )  ↵ sk )^
Revised: Based on the Problem Formulation, a model m sk at stage s can only be activated if I ( q sk ( x t ) ≥ α sk )

Original:  k  1 j =

Revised: ∧ k − 1 j =

Original: I ( q sj ( x t ) <
Revised: I ( q sj ( x t ) <

Original: ↵ sj ) ⌘ 1 .
Revised: α sj ) ≡ 1 .

Original: Now we write the conditional probability of being at stage- s given input x t asPr ( y s = 1 | x t ; m casc s ) =
Revised: Now we write the conditional probability of being at stage- s given input x t as

Original: P
Revised: [Equation]

Original: K s k
Revised: where m sk ( x t )

Original: [-=
-]
Revised: [-=
-]=

Original: I ( q sk ( x t ) ↵ sk ) · Q k  1 j =1 I ( q sj ( x t ) < ↵ sj ) · m sk ( x t ) , where m sk ( x t ) = Pr  y s = 1 | x t ; m sk  .
Revised: Pr (cid:0) y s = 1 | x t ; m sk (cid:1) .

Original: The onestage loss function in Hard-Gating is deﬁned as the negative log-likelihood loss L s nll  m casc s ; D s  =
Revised: The one-stage loss function in Hard-Gating is defined as the negative log-likelihood loss L s nll (cid:0) m casc s ; D s (cid:1) = − (cid:80) N s i =1 (cid:80)

Original: P
Revised: [Equation]

Original: N s i = 
Revised: = [Equation]

Original: P
Revised: sit

Original: T
Revised: [Equation]

Original: i t =1 y sit · log p sit +(1  y sit ) ·
Revised: sit [Equation] [Equation]

Original:  log(1  p sit ) , where p sit =
Revised: y sit ) · log(1 − p sit ) , where p s it =

Original: Pr ( y s = 1 | x it ; m casc s ) , y sit = 1 onlyif y si = 1 and t 2 [ t 1 i   t , t 1 i ] for some  t time-steps we wish to early detect the next stage s + 1 .
Revised: Pr ( y s = 1 | x it ; m casc s ) , y s it = 1 only if y s i = 1 and t ∈ [ t 1 i − δ t , t 1 i ] for some δ t time-steps we wish to early detect the next stage s + 1 .

Original: Algorithm 1 in Appendix 5.1 describes a bottom-up grid search algorithm for for learning hard-gatingparameters.
Revised: Algorithm 1 in Appendix 5.1 describes a bottom-up grid search algorithm for for learning hard-gating parameters.

Original: We sort the model list M s = { m s 1 , · · · m sK s } in a monotonically increasing order185 w.r.t AUC and cost which discards any sub-optimal models.
Revised: We sort the model list M s = { m s 1 , · · · m sK s } in a monotonically increasing order w.r.t AUC and cost which discards any sub-optimal models.

Original: For any stage s , Algorithm 1 starts byassigning all the N s samples for that stage to the ﬁrst model m s 1 , and performing a grid search onthe gate parameters ↵ sk ’s for level k
Revised: For any stage s , Algorithm 1 starts by assigning all the N s samples for that stage to the first model m s 1 , and performing a grid search on the gate parameters α sk ’s for level k

Original: = 1 to K s .
Revised: = 1 to K s .

Original: It gradually assigns IDK samples to the next level’smodel in the list until the cost exceeds the user-deﬁned budget c s for that stage.
Revised: It gradually assigns IDK samples to the next level’s model in the list until the cost exceeds the user-defined budget c s for that stage.

Original: In each iteration ofsearching the cutoff ↵ sk , we set an upper bound maxA on the maximum searching value to avoidover-upgrading.
Revised: In each iteration of searching the cutoff α sk , we set an upper bound maxA on the maximum searching value to avoid over-upgrading.

Original: Without setting this bound, the algorithm could overﬁt and put all the samples intothe IDK class, then assign them to the next level’s model.
Revised: Without setting this bound, the algorithm could overfit and put all the samples into the IDK class, then assign them to the next level’s model.

Original: This consumes the cost quota quickly,and prevents high IDK samples from exploring larger models in the list.
Revised: This consumes the cost quota quickly, and prevents high IDK samples from exploring larger models in the list.

Original: 3.2.2 Soft-gating Training Algorithm
Revised: 3.2.2 Soft-gating Training Algorithm

Original: In contrast to grid search on the thresholds, we propose a soft-gating algorithm formulating anobjective function that can be efﬁciently solved using gradient descent algorithms.
Revised: In contrast to grid search on the thresholds, we propose a soft-gating algorithm formulating an objective function that can be efficiently solved using gradient descent algorithms.

Original: In this algorithm,we deﬁne the gate function G IDK to be a ReLU function parameterized by a pair of coefﬁcients( a sk , b sk ) s.t.
Revised: In this algorithm, we define the gate function G IDK to be a ReLU function parameterized by a pair of coefficients ( a sk , b sk ) s.t.

Original: the gate is only activated if the linear product a sk · q sk ( x t )  b sk > 0 .
Revised: the gate is only activated if the linear product a sk · q sk ( x t ) − b sk > 0 .

Original: Formally wedeﬁne Soft-gating asSoft
Revised: Formally we define Soft-gating as

Original: [Equation]

Original: gating

Original: [Equation]
Revised: [Equation]

Original: Therefore the conditional probability of being at the stage-
Revised: Therefore the conditional probability of being at the stage-

Original: s given input x t is deﬁned as a mixture ofthe K s models available for stage- s prediction:
Revised: s given input x t is defined as a mixture of the K s models available for stage- s prediction:

Original: Pr( y s = 1 | x t ; m casc ) =
Revised: Pr

Original: [Equation]
Revised: [Equation]

Original: Now, the negative log-likelihood loss L s nll becomes solvable using gradient descent algorithms.
Revised: Now, the negative log-likelihood loss L s nll becomes solvable using gradient descent algorithms.

Original: However, the normalization term in the mixture of experts requires running all the candidates modelsin the zoo, which conﬂicts with our cost-saving goal.
Revised: However, the normalization term in the mixture of experts requires running all the candidates models in the zoo, which conflicts with our cost-saving goal.

Original: Therefore, we propose using a Dirichlet
Revised: Therefore, we propose using a Dirichlet

Original: Knowledge Distillation (DKD) for training a small surrogate model for each stage to quantify theprediction conﬁdence for each model in the zoo.
Revised: Knowledge Distillation (DKD) for training a small surrogate model for each stage to quantify the prediction confidence for each model in the zoo.

Original: Then we replace the q sj ( x t ) ’s with their estimationsˆ q sj ( x t ) ’s when selecting the expert models to infer with when real predictions are made.
Revised: Then we replace the q sj ( x t ) ’s with their estimations ˆ q sj ( x t ) ’s when selecting the expert models to infer with when real predictions are made.

Original: The smallerdistilled model only needs to be run once per query, requiring much less cost than running all themodels.
Revised: The smaller distilled model only needs to be run once per query, requiring much less cost than running all the models.

Original: In our experiment, we utilize the ﬁrst model m s 1 from each stage, take the embedding ofh s 1 ( x t ) prior to the last activation layers in m s 1 and feed it into a 4-layer K s -head Multi-Layer
Revised: In our experiment, we utilize the first model m s 1 from each stage, take the embedding of h s 1 ( x t ) prior to the last activation layers in m s 1 and feed it into a 4-layer K s -head Multi-Layer

Original: Perceptron (MLP) that is the distilled model.
Revised: Perceptron (MLP) that is the distilled model.

Original: The idea of DKD is to ﬁrst posit a Dirichlet prior distribution over the parameters ⇡ characterizingthe output categorical distribution (e.g., binomial in our problem setup) predicted from a model asfollows:
Revised: The idea of DKD is to first posit a Dirichlet prior distribution over the parameters π characterizing the output categorical distribution (e.g., binomial in our problem setup) predicted from a model as follows: Pr [Equation] t

Original: [Equation]
Revised: [Equation]

Original:  surrogate priornetwork f is ﬁt to generate the concentration ↵ sk of the Dirichlet distribution.
Revised: sk ) , A surrogate prior network f is fit to generate the concentration α sk of the Dirichlet distribution.

Original: If the learnt concentration parameters yield a ﬂat prior distribution, then it means high uncertainty in the model prediction; if they yield a sharp prior distribution, it means low uncertainty.
Revised: If the learnt concentration parameters yield a flat prior distribution, then it means high uncertainty in the model prediction; if they yield a sharp prior distribution, it means low uncertainty.

Original: Then an estimation ˆ q sk ( x t ) can be computed from the expected predictive probabilityˆ p sk ( x t ) =
Revised: Then an estimation ˆ q sk ( x t ) can be computed from the expected predictive probability ˆ p sk ( x t ) =

Original: E ⇡ ⇠ Dir ( ⇡ | ↵ sk ) [ ⇡ 1 ] =
Revised: E π ∼ Dir ( π | α sk ) [ π 1 ] =

Original: ↵ sk, 1 / ( ↵ sk, 0{+
+} + ↵ sk, 1 ) .
Revised: α sk, 1 / ( α sk, {+
+} + α sk, 1 ) .

Original: For training the DKD model, which219 is a K s -head
Revised: For training the DKD model, which is a K s -head

Original: MLP per each stage in this paper, we deﬁne the loss function for head k as a
Revised: MLP per each stage in this paper, we define the loss function for head k as a

Original: Kullback-Leibler (KL) divergence between the prior distribution and empirical observed distribution: L ( ↵
Revised: Kullback-Leibler (KL) divergence between the prior distribution and empirical observed distribution: L ( α sk ) = (cid:80) N s i

Original: [Equation]
Revised: =1 (cid:80) T

Original:  x it )  , where p sk ( x it ) are the real predictedprobabilities produced from model m sk on input x it .
Revised: i t =1 KL (cid:0) Dir ( π | α sk ) || p sk ( x it ) (cid:1) , where p sk ( x it ) are the real predicted probabilities produced from model m sk on input x it .

Original: Additionally, we also add the cross-entropyloss as an auxiliary loss when training the DKD.
Revised: Additionally, we also add the cross-entropy loss as an auxiliary loss when training the DKD.

Original: Once the distilled model is trained, we replace the q sj ( x t ) ’s with the estimated values from the model.
Revised: Once the distilled model is trained, we replace the q sj ( x t ) ’s with the estimated values from the model.

Original: Then for cost constraining, we ensure the selectedmodels do not exceed user-imposed cost constraint by enforcing sparse gating weights over the modelchoices.
Revised: Then for cost constraining, we ensure the selected models do not exceed user-imposed cost constraint by enforcing sparse gating weights over the model choices.

Original: Therefore, we reformulate the objective function in Hard-Gating Sub-Objective 1 as thefollowing Lagrangian function min G IDK L s nll ( m casc ; D s ) 
Revised: Therefore, we reformulate the objective function in Hard-Gating Sub-Objective 1 as the following Lagrangian function min G IDK L s nll ( m casc ; D s ) + λ L s cost + µ L s sparse , where the second term

Original: +  L s cost + µ L s sparse , where the second term{+cost constraint, controlled by λ > 
+}
Revised: L s cost takes the {+cost constraint, controlled by λ > 
+}

Original: L s cost takes the cost constraint, controlled by  > 0 and the third term L s sparse imposes further sparseregularization on the gating weights w.r.t their L 1 norms, controlled by µ > 0 . Formally, we writethe last two terms as L s cost = 
Revised: and the third term L s sparse imposes further sparse regularization on the gating weights w.r.t their L 1 norms, controlled by µ > 0 . Formally, we write the last two terms as L s cost = (cid:0) max(0 , cost ( m casc s ) − c s ) (cid:1) 2 and L s sparse = (cid:80) i,t ∥ G IDK (ˆ q sk ( x it ))

Original: [Equation]
Revised: ∥ 1 .

Original: Now we use stochastic gradient descent algorithms to learn the optimal a s and b s minimizing theabove loss.
Revised: Now we use stochastic gradient descent algorithms to learn the optimal a s and b s minimizing the above loss.

Original: Algorithm 2 in Appendix 5.1 summarizes the soft-gating training algorithm.
Revised: Algorithm 2 in Appendix 5.1 summarizes the soft-gating training algorithm.

Original: 3.2.3 Overall Training Algorithm
Revised: 3.2.3 Overall Training Algorithm

Original: To complete the overall training algorithm, we need to learn the optimal gate parameter G ICK 1 in
Revised: To complete the overall training algorithm, we need to learn the optimal gate parameter G ICK 1 in

Original: Sub-Objective 2.
Revised: Sub-Objective 2.

Original: Given the model m sk picked by one-stage IDK cascade (either using hard-gatingor soft-gating algorithms) for a data example x t at stage s , we grid search for the optimal thresholds✓ sk ’s on the predictive probabilities s.t.
Revised: Given the model m sk picked by one-stage IDK cascade (either using hard-gating or soft-gating algorithms) for a data example x t at stage s , we grid search for the optimal thresholds θ sk ’s on the predictive probabilities s.t.

Original: the type I error on class ICK 1 and type II error on ICK 0are both minimized.
Revised: the type I error on class ICK 1 and type II error on ICKare both minimized.

Original: Several existing methods (Liu, 2012; Perkins and Schisterman, 2006; Unal,2017; Miller and Siegmund, 1982) have been proposed for minimizing both type I error and type
Revised: Several existing methods (Liu, 2012; Perkins and Schisterman, 2006; Unal, 2017; Miller and Siegmund, 1982) have been proposed for minimizing both type I error and type

Original: II error in various ways, we pick the Closet-to-(0,1) (Perkins and Schisterman, 2006) method thatﬁnds the optimal threshold achieving the most left upper corner in the ROC curve.
Revised: II error in various ways, we pick the Closet-to-(0,1) (Perkins and Schisterman, 2006) method that finds the optimal threshold achieving the most left upper corner in the ROC curve.

Original: Finally, Algorithm3 in Appendix 5.1 gives the end-to-end training Algorithm for UnfoldML , where we can use either
Revised: Finally, Algorithm 3 in Appendix 5.1 gives the end-to-end training Algorithm for UnfoldML , where we can use either

Original: Algorithm 1 or 2 to learn the optimal gating policy.
Revised: Algorithm 1 or 2 to learn the optimal gating policy.

Original:  Experiments
Revised:  Experiments

Original: We evaluate UnfoldML on two real-world tasks.
Revised: We evaluate UnfoldML on two real-world tasks.

Original: The ﬁrst task is to predict if and when a patient whowas newly admitted into the Intensive Care Unit (ICU) of a hospital will develop sepsis (Stage-1),which can then progress into septic shock (Stage-2).
Revised: The first task is to predict if and when a patient who was newly admitted into the Intensive Care Unit (ICU) of a hospital will develop sepsis (Stage-1), which can then progress into septic shock (Stage-2).

Original: The second task is detection of the subcategoryof-interest in a label hierarchy that ﬁlters out queries that are not of interest in Stage-1 and then reﬁnesthe predictions into ﬁne classes in Stage-2.
Revised: The second task is detection of the subcategoryof-interest in a label hierarchy that filters out queries that are not of interest in Stage-1 and then refines the predictions into fine classes in Stage-2.

Original: 4.1 Task 1: Sepsis-Septic Shock prediction
Revised: 4.1 Task 1: Sepsis-Septic Shock prediction

Original: We use MIMIC-III Critical Care Database (Johnson et al., 2016).
Revised: We use MIMIC-III Critical Care Database (Johnson et al., 2016).

Original: The database consists of deidentiﬁedhealth records from over 50 , 000 critically ill patients who stayed in the ICUs of the Beth Israel
Revised: The database consists of deidentified health records from over 50 , 000 critically ill patients who stayed in the ICUs of the Beth Israel

Original: Deaconess Medical Center between 2001 and 2012.
Revised: Deaconess Medical Center between 2001 and 2012.

Original: We detail our data preparation in Appendix [-5.
-]
Revised: We detail our data preparation in Appendix [-5.
-]5.

Original: The ﬁnal cohort includes a total of 34 , 475 ICU patients, from which 2 , 370 (6 . 8%) presented with
Revised: The final cohort includes a total of 34 , 475 ICU patients, from which 2 , 370 (6 . 8%) presented with

Original: Sepsis, from which a total of 229 (9 . 7%) progressed into Septic Shock.
Revised: Sepsis, from which a total of 229 (9 . 7%) progressed into Septic Shock.

Original: We randomly split our cohortof patient data into a training set ( 70% ), validation set ( 20% ) and test set ( 10% ).
Revised: We randomly split our cohort of patient data into a training set ( 70% ), validation set ( 20% ) and test set ( 10% ).

Original: First, we use thetraining set to train a set of models to formulate a model zoo.
Revised: First, we use the training set to train a set of models to formulate a model zoo.

Original: Next, we use the validation set to trainthe UnfoldML policy.
Revised: Next, we use the validation set to train the UnfoldML policy.

Original: Finally, we use the test set to evaluate performance.
Revised: Finally, we use the test set to evaluate performance.

Original: 4.1.1 Experimental Setup
Revised: 4.1.1 Experimental Setup

Original: Model Zoo.
Revised: Model Zoo.

Original: We construct our model zoo by training two sets of binary classiﬁers for Stage-1 (Sepsis)and for Stage-2 (Septic Shock) using a variety of architectures and features.
Revised: We construct our model zoo by training two sets of binary classifiers for Stage-1 (Sepsis) and for Stage-2 (Septic Shock) using a variety of architectures and features.

Original: We select CPU-basedmodels such as Logistic Regression, Decision Tree and Random Forest, and GPU-based models suchas LSTM (Long Short-Term Memory) that is thus far the state-of-art approach in the early detectionof Sepsis and Septic Shock (Fagerström et al., 2019; Liu et al., 2019).
Revised: We select CPU-based models such as Logistic Regression, Decision Tree and Random Forest, and GPU-based models such as LSTM (Long Short-Term Memory) that is thus far the state-of-art approach in the early detection of Sepsis and Septic Shock (Fagerström et al., 2019; Liu et al., 2019).

Original: For LSTM, we vary the hiddensize ranging from 100 to 400, the number of layers from 1 to 4.
Revised: For LSTM, we vary the hidden size ranging from 100 to 400, the number of layers from 1 to 4.

Original: We also vary the window of patientdata we input over 1, 6, and 12 hours.
Revised: We also vary the window of patient data we input over 1, 6, and 12 hours.

Original: For each stage and model architecture we train models withdifferent combinations of feature sets based on the collection modality.
Revised: For each stage and model architecture we train models with different combinations of feature sets based on the collection modality.

Original: All models start with withbasic demographic features and vital signs, and are then extended with additional features includingresults from lab tests, beside monitoring, and medication/IV treatments. 
Revised: All models start with with basic demographic features and vital signs, and are then extended with additional features including results from lab tests, beside monitoring, and medication/IV treatments. Confidence Measures We consider four choices for measuring the confidence q of a model prediction.

Original: Spatio-Temporal Cost.{+Given p , the model’s predictive probability of output Y being 1 , we define • Max probability: max (cid:0) 1 − p, p (cid:1) , • Entropy: (cid:0) p · log p + (1 − p ) · log(1 − p ) (cid:1) , • Entropy of expected: − α 
 α 
 + α 1 · (cid:0) ψ ( α 0 )+}
Revised: {+Given p , the model’s predictive probability of output Y being 1 , we define • Max probability: max (cid:0) 1 − p, p (cid:1) , • Entropy: (cid:0) p · log p + (1 − p ) · log(1 − p ) (cid:1) , • Entropy of expected: − α 
 α 
 + α 1 · (cid:0) ψ ( α 0 )+}

Revised: − ψ ( α 1 ) (cid:1) + ψ ( α 0 + α 1 ) , • Mutual Information: Entropy − Entropy of expected,

Original:  Given the trained model zoo, we proﬁle spatio-temporal costs for each of themodels.
Revised: Spatio-Temporal Cost. Given the trained model zoo, we profile spatio-temporal costs for each of the models.

Original: The spatio-temporal cost is the total time spent in each hardware due to inference/forwardpass calls of the models (temporal cost) multiplied with the hardware’s cost per unit time (spatial[Figure] cost).
Revised: The spatio-temporal cost is the total time spent in each hardware due to inference/forwardpass calls of the models (temporal cost) multiplied with the hardware’s cost per unit time (spatial cost).

Original: This cost serves as a proxy for the real dollar cost as it is the basis for pricing models in cloudofferings such as AWS On-Demand, Lambda, or Spot.
Revised: This cost serves as a proxy for the real dollar cost as it is the basis for pricing models in cloud offerings such as AWS On-Demand, Lambda, or Spot.

Original: Baseline.
Revised: Baseline.

Original: To our knowledge, UnfoldML is the ﬁrst system to provide 2-dimensional cascadingpredictions.
Revised: To our knowledge, UnfoldML is the first system to provide 2-dimensional cascading predictions.

Original: A reasonable baseline for our method are models that frame the multi-stage sequentialtask as a multi-class classiﬁcation task by ignoring the happen-before relationship between thestages.
Revised: A reasonable baseline for our method are models that frame the multi-stage sequential task as a multi-class classification task by ignoring the happen-before relationship between the stages.

Original: The baseline’s goal is to classifying patients into one of the three classes: Non-Septic, Septic,and Septic Shock.
Revised: The baseline’s goal is to classifying patients into one of the three classes: Non-Septic, Septic, and Septic Shock.

Original: We use a state-of-the-art LSTM which is widely used in prediction of clinicaltime-series.
Revised: We use a state-of-the-art LSTM which is widely used in prediction of clinical time-series.

Original: We evaluate performance LSTM’s across a variety of spatio-temporal costs by changingthe number of layers and hidden sizes in the architecture.
Revised: We evaluate performance LSTM’s across a variety of spatio-temporal costs by changing the number of layers and hidden sizes in the architecture.

Original: • Multi-class LSTM :
Revised: • Multi-class LSTM :

Original: One uniﬁed multi-class model works end to end forpredicting the
Revised: One unified multi-class model works end to end for

Original:  multi classes
Revised: predicting the multi classes.

Original: Prior work (Wang et al., 2017) implements a 1-D prediction cascade on IDK classes.
Revised: Prior work Wang et al., 2017) implements a 1-D prediction cascade on IDK classes.

Original: It uses a similarhard-gating algorithm to learn hard thresholds on prediction entropy for making decisions if thesystem should cascade to costlier models.
Revised: It uses a similar hard-gating algorithm to learn hard thresholds on prediction entropy for making decisions if the system should cascade to costlier models.

Original: So we reduce UnfoldML to a 1-D pipeline by removing itshorizontal ICK transition and compare the soft-gating algorithm with this baseline.
Revised: So we reduce UnfoldML to a 1-D pipeline by removing its horizontal ICK transition and compare the soft-gating algorithm with this baseline.

Original: We evaluate themon the two single-stage classiﬁcation tasks on sepsis and septic shock predictions respectively.• Single-Stage binary classiﬁers : CPU-based models such as Logistic Regression, Decision Tree,Random Forest and TREWScore (Henry et al., 2015), a cox proportional hazards model that iswell-known in early detection of septic shock; GPU-based models such as LSTMs.• IDK-cascade (Wang et al., 2017): Prior work of 1-D prediction cascade on IDK classes using ahard-gating like algorithm.
Revised: We evaluate them on the two single-stage classification tasks on sepsis and septic shock predictions respectively. • Single-Stage binary classifiers : CPU-based models such as Logistic Regression, Decision Tree, Random Forest and TREWScore (Henry et al., 2015), a cox proportional hazards model that is well-known in early detection of septic shock; GPU-based models such as LSTMs. • IDK-cascade (Wang et al., 2017): Prior work of 1-D prediction cascade on IDK classes using a hard-gating like algorithm.

Original: 4.1.2 Evaluation and Results
Revised: 4.1.2 Evaluation and Results

Original: End-to-End Classiﬁcation Performance.
Revised: End-to-End Classification Performance.

Original: To evaluate the prediction performance on our multi-stagetask, we treat UnfoldML as a multi-class classiﬁer which predicts one of the same three classes asdeﬁned in Baseline.
Revised: To evaluate the prediction performance on our multi-stage task, we treat UnfoldML as a multi-class classifier which predicts one of the same three classes as defined in Baseline.

Original: UnfoldML generates predictions at time-step for every patient, we take themaximum of the predictive probabilities along the prediction horizon for each patient and normalizethem with a sum of 1.
Revised: UnfoldML generates predictions at time-step for every patient, we take the maximum of the predictive probabilities along the prediction horizon for each patient and normalize them with a sum of 1.

Original: We compute the multi-class ROC AUC scores by averaging the pairwise ROC
Revised: We compute the multi-class ROC AUC scores by averaging the pairwise ROC

Original: AUCs (known as one-vs-one) of each classes.
Revised: AUCs (known as one-vs-one) of each classes.

Original: Better Cost-AUC tradeoff.
Revised: Better Cost-AUC tradeoff.

Original: We evaluate the model performance in trading off between the end-to-end
Revised: We evaluate the model performance in trading off between the end-to-end

Original: AUC and spatio-temporal cost as follows:
Revised: AUC and spatio-temporal cost as follows:

Original: we vary the user-deﬁned cost constraint c in Eq. 1, andtrain UnfoldML repeatedly while searching the trade off in a 2-D space.
Revised: we vary the user-defined cost constraint c in Eq. 1, and train UnfoldML repeatedly while searching the trade off in a 2D space.

Original: Each run contributes a pointin the scatter plot of Figure 2b.
Revised: Each run contributes a point in the scatter plot of Figure 2b.

Original: We compute the convex hull of the searched points in the set anddemonstrate the resulting area of the hull in Figure 2a.
Revised: We compute the convex hull of the searched points in the set and demonstrate the resulting area of the hull in Figure 2a.

Original: We train our baseline method repeatedly byvarying the LSTM architecture and plot in the same way as in Figure 2a.
Revised: We train our baseline method repeatedly by varying the LSTM architecture and plot in the same way as in Figure 2a.

Original: In comparison, UnfoldML -S( soft-gating ) searches signiﬁcantly higher AUC regions than the multi-class LSTM baseline andperforms more consistently than the hard-gating algorithm.
Revised: In comparison, UnfoldML -S ( soft-gating ) searches significantly higher AUC regions than the multi-class LSTM baseline and performs more consistently than the hard-gating algorithm.

Original: To quantify the trade-off evaluation in
Revised: To quantify the trade-off evaluation in

Original: Figure 2a, we calculate the area of the convex hull searched by each method (Table 1). 
Revised: Figure 2a, we calculate the area of the convex hull searched by each method (Table 1). [Figure] nd AU

Revised: [Figure]

Original:  UnfoldML -S presents the highest score comparing to the baseline with UnfoldML
Revised: [Figure] [Figure] [Figure] UnfoldML -S presents the highest score comparing to the baseline with UnfoldML

Original: -H. In addition,we pick several critical points from the searched convex hull in Figure 2b and present them in
Revised: -H. In addition, we pick several critical points from the searched convex hull in Figure 2b and present them in

Original: [Table]

Original: Table 2.
Revised: Table 2.

Original: The recommended conﬁguration for systems that effectively trade off spatio-temporalcost and end-to-end AUC is to select models along the frontier line drawn through the points thatform the top-left corner of the hull.
Revised: The recommended configuration for systems that effectively trade off spatio-temporal cost and end-to-end AUC is to select models along the frontier line drawn through the points that form the top-left corner of the hull.

Original: As shown in Table 2, an optimal conﬁguration for our pipeline
Revised: As shown in Table 2, an optimal configuration for our pipeline

Original: UnfoldML -S (Rec.) can achieve comparable end-to-end AUC score ( 88 . 9% ) with the most accuratebaseline Multi-Class LSTM-A ( 88 . 8% ), while operating at a 19.6X spatio-temporal cost reduction.
Revised: UnfoldML -S (Rec.) can achieve comparable end-to-end AUC score ( 88 . 9% ) with the most accurate baseline Multi-Class LSTM-A ( 88 . 8% ), while operating at a 19.6X spatio-temporal cost reduction.

Original: In comparison with UnfoldML -A, which naively uses the most accurate single-stage models in eachstage, UnfoldML -S outperforms the end-to-end AUC while providing a 22.7X reduction in spatiotemporal cost.
Revised: In comparison with UnfoldML -A, which naively uses the most accurate single-stage models in each stage, UnfoldML -S outperforms the end-to-end AUC while providing a 22.7X reduction in spatiotemporal cost.

Original: Alternatively, UnfoldML -S (Rec.) outperforms Multi-class LSTM
Revised: Alternatively, UnfoldML -S (Rec.) outperforms Multi-class LSTM

Original: (Rec.)
Revised: (Rec.)

Original: at an AUCgain of 5 . 3% with only 1.3X more cost.
Revised: at an AUC gain of 5 . 3% with only 1.3X more cost.

Original: The improvement in end-to-end AUC scores from UnfoldMLis a result of the dynamic nature of inference pipeline, which selects the optimal pathway for eachpatient.
Revised: The improvement in end-to-end AUC scores from UnfoldML is a result of the dynamic nature of inference pipeline, which selects the optimal pathway for each patient.

Original: The savings in spatio-temporal cost can be attributed to the usage of low cost models whenthey are conﬁdent enough, in contrast, baselines ( UnfoldML -A, Multi-Class LSTM-A) always usethe most accurate model.[Table] [Table] [Table]
Revised: The savings in spatio-temporal cost can be attributed to the usage of low cost models when they are confident enough, in contrast, baselines ( UnfoldML -A, Multi-Class LSTM-A) always use the most accurate model.

Original: [Table]

Original: Single-Stage Performance.
Revised: Single-Stage Performance.

Original: We report single-stage performance on Sepsis and Septic Shock predictions in Table 3.
Revised: We report single-stage performance on Sepsis and Septic Shock predictions in Table 3.

Original: For both predictions, UnfoldML achieves a signiﬁcant reduction in costs withonly a marginal loss of AUC compared to the highest accuracy baseline LSTM-A — 32.3x lowerspatio-temporal cost with 1.7% lower AUC for Sepsis prediction and 26.8x lower costs with 0.7%lower AUC for Septic Shock prediction.
Revised: For both predictions, UnfoldML achieves a significant reduction in costs with only a marginal loss of AUC compared to the highest accuracy baseline LSTM-A — 32.3x lower spatio-temporal cost with 1.7% lower AUC for Sepsis prediction and 26.8x lower costs with 0.7% lower AUC for Septic Shock prediction.

Original: The baseline IDK-Cascade can achieve comparable AUC as
Revised: The baseline IDK-Cascade can achieve comparable AUC as

Original: UnfoldML -S, but requires 3.4x higher costs.
Revised: UnfoldML -S, but requires 3.4x higher costs.

Original: Better Early-hour Prediction.
Revised: Better Early-hour Prediction.

Original: Table 2 also reports the average early prediction, showing UnfoldMLS predicts septic shock earlier than all other methods.
Revised: Table 2 also reports the average early prediction, showing UnfoldMLS predicts septic shock earlier than all other methods.

Original: It predicts 2.1 hrs prior to the strongest baseline
Revised: It predicts 2.1 hrs prior to the strongest baseline

Original: Mutli-class LSTM-A. It beneﬁts from the multi-stage cascaded prediction in UnfoldML such thatthe system can exit early from the sepsis stage once it is in the ICK 1 class and proceed to shockprediction before sepsis is truly diagnosed.
Revised: Mutli-class LSTM-A. It benefits from the multi-stage cascaded prediction in UnfoldML such that the system can exit early from the sepsis stage once it is in the ICK 1 class and proceed to shock prediction before sepsis is truly diagnosed.

Original: 4.2 Task 2: Subcategory Classiﬁcation
Revised: 4.2 Task 2: Subcategory Classification

Original: Our second experiment is a computer vision task demonstrating the generalizability of UnfoldML  .
Revised: To demonstrate the generalizability of UnfoldML we conduct experiments on a computer vision task where we wish to detect a subcategory-of-interest .

Original: Our goal identiﬁcation of images from a  subcategory in a dataset with  hierarchical  label structure.
Revised: In the subcategory-of-interest task, we wish to accurately predict if an image falls within a specific subcategory of a dataset with classes that form a hierarchical structure, while still ensuring our spatio-temporal costs remain within the user-defined budget. In this task, each query is an image from which UnfoldML is asked to predict if the image belongs to a subcategory in the label hierarchy we are interested in.

Original: For our experiment we use the  CIFAR-100 dataset (Krizhevsky et al., 2009), and deﬁne two distinctsubcategory-of-interest tasks based on  real-world  computer vision systems.
Revised: For our experiment using the [Table] [Table] CIFAR-100 dataset (Krizhevsky et al., 2009), we define two separate subcategory-of-interest tasks based on the real-world system applications of computer vision systems.

Original: We want to identifyimages with the subcategory fromfrom one of two  categories: ‘people’ or ‘vehicles .
Revised: We want to identify images of a specific subcategory from one of two chosen categories: ‘people’ (baby, boy, girl, man, woman) and ‘vehicles’ (bicycle, bus, motorcycle, pickup truck, train).

Original: CIFAR-100consists of 60,000 32x32 colour images with a total of20 ‘coarse’ labels and 100 ‘ﬁne’ labels.
Revised: CIFAR-100 consists of 60,000 32x colour images in total of 100 ‘fine’ classes.

Original: For each ‘ﬁne’ class there are 500 training  and 100 testing images.
Revised: There are 500 training images and 100 testing images per class.

Revised: In addition, each image also is assigned a ‘coarse’ label indicating which category (such as people, vehicles, trees, etc) it belongs to.

Original: In our experiments Stage-1 predictsthe ‘coarse’ label, and Stage-2 predicts the ‘ﬁne’ label.
Revised: In our experiment, we do a multi-class classification on the ’coarse-granularity’ labels in Stage-1 to detect the category, then we do multi-class classification on the original ‘fine-granularity’ CIFAR-100 labels in the second step to identify the subcategory-of-interest .

Original: 4.2.1 Experimental Setup
Revised: 4.2.1 Experimental Setup

Original: Our model zoo is made of WideResNets (WRN) using Sharpness-Aware Minimization (Foret et al. ,2020) which is cited as having state-of-the-art performance on the CIFAR-100 image recognition task.
Revised: Our model zoo is made of WideResNets (WRN) using Sharpness-Aware Minimization (Foret et al., 2020) which is cited as having state-of-the-art performance on the CIFAR-100 image recognition task.

Original: For each stage we train models with widths from [2 , 4 , 6 , 8 , 10] and depths from [16 , 22 , 24 , 28] .
Revised: For each stage we train models with widths from [2 , 4 , 6 , 8 , 10] and depths from [16 , 22 , 24 , 28] .

Original: For
Revised: For

Original: Stage-1 we train  multi-class  models using  full dataset with ‘coarse’ labels as the target.
Revised: Stage-1 we train K 1 number of multi-class classification models using the full dataset of CIFARwith coarse-granularity labels as our target subcategories, resulting in a total of 20 classes.

Revised: We train models for Stage-1 using the full 32x32 image and as well as random crops of size 24x24 or 16x16.

Original: For Stage-2we target the ‘ﬁne’ subcategory that happens-after the ’coarse‘ subcategory from Stage-1.
Revised: For Stage-2 we train K 2 multi-class models to identify the next subcategory that happens-after the coarse subcategory from Stage-1.

Original: We train  Stage-2 models  using data from the coarse-granularity label.
Revised: We train our Stage-2 models only using data with the coarsegranularity label for our subcategory in Stage-1.

Original: For Stage-1 we use the 32x32 image,or a random crop of size 24x24 or 16x16.

Original: For Stage-2 we train  using the  32x32 image or randomcrop of 24x24.
Revised: We train Stage-2 models using the full 32x32 image or random crop of 24x24.

Original: We proﬁle each model to determine the number of Multiply-Accumulate Operations(MACs) (using toolkit 1 ) which we use as proxy measure of  spatio-temporal cost.
Revised: We profile each model to determine the number of Multiply-Accumulate Operations (MACs) (using toolkit 3 ) which we use as proxy measure of our spatio-temporal cost.

Original: 4.2.2 Evaluation and Results
Revised: 4.2.2 Evaluation and Results

Original: [Table] We compare UnfoldML against a baseline multiclass WRN classiﬁer which we train for directlypredicting labels at the ﬁne-granularity.
Revised: We compare UnfoldML against a baseline using a multi-class WRN classifier which we train for directly predicting labels at the fine-granularity.

Original: We trainour baseline with a width factor of 10 and adepth of 28 using the full-sized 32x32 image.
Revised: We train our baseline with a width factor of 10 and a depth of 28 using the full-sized 32x32 image.

Original: We evaluate accuracy on our subcategory-ofinterest task and report our results in Table 4.
Revised: We evaluate accuracy on our subcategory-of-interest task and report our results in Table 4.

Original: Our results show that UnfoldML is able to ﬁnd the optimal gating policy for the pipeline and achievea spatio-temporal cost-savings of 6.9X with 1% reduction in accuracy for ‘people’ subcategoryclassiﬁcation; we achieve a cost-savings of 4.7X with 0.4%  reduction in accuracy for ‘vehicle’subcategory-of-interest identiﬁcation.
Revised: Our results show that UnfoldML is able to find the optimal gating policy for the pipeline and achieve a spatio-temporal cost-savings of 6.9X with 1% reduction in accuracy for ‘people’ subcategory classification; we achieve a cost-savings of 4.7X with 0.4% [Table] reduction in accuracy for ‘vehicle’ subcategory-of-interest identification. Thus, UnfoldML achieves a savings in spatio-temporal cost in for this task without compromising the overall accuracy.

Original: An alternative method is Classiﬁcation with Costly Features (CwCF) (Janisch et al., 2019), whichnavigates trade-off space using Deep Q-Learning while varying cost parameters which restrictsusing features when training models. 
Revised: UnfoldML can accomplish this because when it is confident that queries do not belong to our subcategory-ofinterest , we are able to early exit ( ICK 0 ) the query from the prediction pipeline. When UnfoldML is not confident ( IDK ), it upgrades queries to costlier and more accurate models. Only when UnfoldML can confidently predict the subcategory ( ICK 1 ) in Stage-1 will it transition the query to

Original: CWCF requires a pre-trained cost-unaware classiﬁer, whichis used when no features are restricted.
Revised: Stage-2. A reduced set of the original queries transitions to Stage-2, where UnfoldML will predict the subcategory using only the fine-granularity labels. This results in a significant cost-savings without compromising accuracy. A recent alternative method for trading-off accuracy and cost in image classification is CwCF (Classification with Costly Features)

Revised: (Janisch et al., 2019), which traverses the trade-off space by varying a cost parameter that limits the number of selected features and uses Deep Q-Learning (DQL) model to train a feature-cost-aware classifier.

Revised: In addition, it also includes a pre-trained cost-unaware

Revised: High-Performance Classifier (HPC), which is called when it decides to include all the features.

Original: As a fair baseline against UnfoldML , we used the highestaccuracy WRN  in CwCF.
Revised: In order to establish a fair baseline against UnfoldML , we used the most accurate multi-class WRN for the HPC in CwCF.

Original: We vary the cost parameter and search for  the optimal trade-off point in theaccuracy-cost space. CWCF returns two extreme data points with either poor or high accuracy as the WRN-A shows in Table 4.
Revised: We vary their cost parameter lambda for searching the optimal trade-off point in the accuracy-cost space, however, it only returns two extreme data points that gives either low accuracy of 1% with only < 10 features or high accuracy as the WRN-A shows in Table 4 with all the features.

Original: We observed similar failure in traversing the trade-off space for CIFAR-10in Figure 4(c) of the original paper, and therefore we do not include CWCF as a baseline for this task.
Revised: We observe similar failure in traversing the trade-off space for CIFAR-10 in Figure 4(c) of their paper, so we do not include this as a baseline for this task.

Original:  Conclusions
Revised:  Conclusion

Original: ML models, including for healthcare applications, are growing exponentially in size and cost ofinference.
Revised: ML models, including for healthcare applications, are growing exponentially in size and cost of inference.

Original: This is problematic for resource constrained hospital environments.
Revised: This is problematic for resource constrained hospital environments.

Original: Costlier, monolithicmodels require expensive hardware, and can not ﬁt on bed-side compute or even on site computeclusters with clinical implications.
Revised: Costlier, monolithic models require expensive hardware, and can not fit on bed-side compute or even on site compute clusters with clinical implications.

Original: UnfoldML proposes a set of mechanisms and policies to address thegrowing cost of monolithic classiﬁers for healthcare applications.
Revised: UnfoldML proposes a set of mechanisms and policies to address the growing cost of monolithic classifiers for healthcare applications.

Original: It implements a query propagationmechanism that “unfolds” a monolithic multi-class classiﬁer into a sequence of single-class classiﬁers,each with its own cascade progressively more complex models.
Revised: It implements a query propagation mechanism that “unfolds” a monolithic multi-class classifier into a sequence of single-class classifiers, each with its own cascade progressively more complex models.

Original: Each query is allowed to (1)conﬁdently exit the pipeline with an ICK 0 , (2) transition horizontally to the next stage in the pipelinewith an ICK 1 , or (3) upgrade vertically to a more complex model within the horizontal stage with an
Revised: Each query is allowed to (1) confidently exit the pipeline with an ICK 0 , (2) transition horizontally to the next stage in the pipeline with an ICK 1 , or (3) upgrade vertically to a more complex model within the horizontal stage with an

Original: IDK .
Revised: IDK .

Original: This mechanism is coupled with a set of policies, such as soft-gating, that set the thresholdsat which the state transitions occur for queries to the system.
Revised: This mechanism is coupled with a set of policies, such as soft-gating, that set the thresholds at which the state transitions occur for queries to the system.

Original: UnfoldML builds on a fundamentalinsight that classes may have a “happens before” relationship between them, which can be leveragedto “unfold” a classiﬁer, leading to savings in spatio-temporal cost (how much resource used for howlong) and clinically signiﬁcant earlier onset prediction.
Revised: UnfoldML builds on a fundamental insight that classes may have a “happens before” relationship between them, which can be leveraged to “unfold” a classifier, leading to savings in spatio-temporal cost (how much resource used for how long) and clinically significant earlier onset prediction.

Original: UnfoldML improves the frontier of optimalityin the cost/accuracy tradeoff space and is able to nearly match (within 0.1%)
Revised: UnfoldML improves the frontier of optimality in the cost/accuracy tradeoff space and is able to nearly match (within 0.1%)

Original: SoTA AUC performancefor septic shock prediction at the 120 th of the baseline cost.
Revised: SoTA AUC performance for septic shock prediction at the 120 th of the baseline cost.

Original: UnfoldML and the application of the“happens before” insight generalizes to computer vision tasks with 5x cost savings gained for a mere0.4% drop in accuracy.
Revised: UnfoldML and the application of the “happens before” insight generalizes to computer vision tasks with 5x cost savings gained for a mere 0.4% drop in accuracy.

Original: Limitations of this work include demonstrations on more than 2 stages tasks. 
Revised: Limitations of this work include demonstrations on more than 2 stages tasks. Acknowledgements

Revised: This material is based upon work supported by the National Science Foundation under Grant Numbers

Revised: NSF IIS-2106961, CAREER IIS-2144338, and CCF-2029004.

Revised: We would also like to acknowledge Dr.

Revised: Kevin Maher and Dr. Alaa Aljiffry of Children’s Healthcare of Atlanta for their medical insights and clinical guidance as well as the Neurips’22 Area Chairs and reviewers for their insightful feedback, which contributed to the improved quality of this paper.

Revised: Disclaimer:

Revised: Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the National Science Foundation.

Revised: Appendix 5.1 Training Algorithms

Revised: Algorithm 1 presents the Hard-Gating Training Algorithm, Algorithm 2 presents the Soft-Gating

Revised: Training Algorithm, and Algorithm 3 summarizes the End-to-End Training Algorithm.

Revised: Algorithm 1 Hard-gating Algorithm for In-Stage IDK Cascade

Revised: Input

Revised: Output

Revised: 2: α ∗ s =

Revised: ModelAssign cost

Revised: (cid:80) i t C m s

Revised: 3: if c s then return α ∗ s 4: end if 5: for k in range ( K s − 1) do search 6:

Revised: Idx 4 k ← ∪ I ( ModelAssign t

Revised: 7: if Idx 4 k is ∅ then break 8: end if 9: minQ ← min Idx 4 k (cid:8) q sk ( x it ) (cid:9) 10: maxQ ← min( maxA, max Idx 4 k (cid:8) q sk ( x it ) (cid:9) ) .

Revised: 11: α ∗ sk ← minQ .

Revised: 12: for α sk in LinSpace ( minQ, maxQ, nBins ) do 13:

Revised: IDK ← ∪ Idx 4 k

Revised: I (cid:0) q sk ( x it ) ∈

Revised: [ α ∗ sk , α sk ) (cid:1) 14: if IDK is not ∅ then 15: if cost + (cid:80) IDK C [ m sk +1 ] − C

Revised: [ m k ] >

Revised: c s then 16:

Revised: ˆ α s ← α ∗ s + [ α ∗ sk ] ; return α ∗ s 17: end if 18: α ∗ sk ← α sk , 19:

Revised: ModelAssign [ IDK ] ← k + 1 , 20: cost + = (cid:80) IDK C [ m sk +1 ] − C [ m k ] 21: end if 22: end for 23: α ∗ s ← α ∗ s

Revised: + [ α ∗ sk ] 24: end for 25: return α ∗ s 26: end procedure

Revised: 5.2 Data preparation

Revised: By following the definition of Sepsis-3 Singer et al.

Revised: (2016), we identify the sepsis onset to be the time when an increase in the Sequential Organ Failure Assessment (SOFA) score of 2 points or more occurs in response to infections.

Revised: We use the Sepsis-3 toolkit 4 to obtain the suspected infection time in patients, and following the process in Seymour et al.

Revised: (2016) to finally label the onset of sepsis.

Revised: We result at a total number of 20 , 009 sepsis patients out of the 52 , 902 adult patients from MIMIC-III database.

Revised: We exclude those patients who stay in ICUs less than 6 hours and also exclude those patients who developed sepsis within the first 6 hours after ICU admission.

Revised: This reduces our cohort to a total of 34 , 475 ICU patient, and only 2 , 370(6 . 8%) out of them are labeled as sepsis (because 88 .

Revised: 1% of sepsis onsets happened within the first 6 hours after ICU admission and are excluded from our study cohort).

Revised: Then according to Singer et al.

Revised: , we identify the onset of septic shock as

Revised: Output

Revised: 1: procedure S OFT G ATING

Revised: ( D s , M s , f M s , c s , C , q ) 2: lr ← 1 e − 1 , e ← 0 , a s ← 1 , b s ← 0 .3: while e < do 4: ˆ q sj ( x t ) ← q (cid:0) f

Revised: ( x t ; m sk )[1] / (cid:80) f x t ; m sk

Revised: (cid:1)

Revised: ▷ distillation 5:

Revised: L s sparse ← (cid:80) i,t,k || G IDK (ˆ q sj ( x t ))

Revised: ||6:

Revised: L a s , b s ← L s nll + λ L s cost + µ L

Revised: s sparse 7:

Revised: Optimize L a s , using

Revised: Reduce lr by factor 0 .

Revised: 5 once learning stagnates.

Revised: 9: e ← e +10: end while 11: return a ∗ s , b ∗ s 12: end procedure

Revised: Input

Revised: Output

Revised: RAINING ( D

Revised: M ) 2:

Revised: Pre-allocate costs c s for each stage s.

Revised: Step 1: Learn in-stage IDK gate parameters.

Revised: 4: for each stage s do 5: α

Revised: ← HardGating( D s

Revised: M s c s

Revised: C q ) 6: or, a

Revised: , b

Revised: ← SoftGating( D s , M s , c s , C, q ) 7: end for 8: 9:

Revised: Step 2: Learn ICK 1 gate parameters.

Revised: 10: for each model m sk do 11: θ sk ← Grid Search for minimizing (cid:113)

Revised: [Equation]

Revised: 12: end for 13: return α ∗ (or a ∗ , b ∗ ), θ ∗ 14: end procedure

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: when a vasopressor is required to maintain a mean arterial pressure (MAP) ≥ 65 mm Hg and serum lactate level > 2 mmol/L ( > 18 mg/dL).

Revised: We result at 229(9 .

Revised: 7%) septic shock patients out of the 2 , 370 sepsis patients.

Revised: For feature generation, we extract 8 patient static characteristics including age, gender, race, height, weight, sepsis onset hour since ICU admission, whether diagnosed diabetes or on a ventilator at ICU admission.

Revised: Then we extract the dynamic features by obtaining the 8 vital signs, 16 lab measurements, 6 vassopressors, continuous replacement therapies (CRRT), ventilation, 2 intravenous fluids in fluid resuscitation, and 5 additional measurements that are recommended for monitoring during sepsis management.

Revised: The 8 vital signs include heart rate, systolic blood pressure, diastolic blood pressure, mean blood pressure, respiration rate, temperature, SpO2 and glucose.

Revised: The 16 lab measurements include Anion gap, Albumin, Bands, Bicarbonate, Bilirubin, Creatinine, Chloride, Glucose, Hematocrit,

Revised: Hemoglobin, Lactate, Platelet, Potassium, PTT, INR, PT, Sodium, BUN and WBC.

Revised: The 6 vasopressors include dobutamine, dopamine, epinephrine, norepinephrine, phenylephrine, and vasopressin.

Revised: The 2 fluids include Crystalloids and Colloids that are recommended in the early management of sepsis Rhodes et al.

Revised: (2017), and particularly fluid resuscitation of bolus ≥ 500 mL is one of the most common treatment for managing septic shock.

Revised: The 5 additional measurements include whether a vasopressor is needed to maintain a mean arterial pressure (MAP) ≥

Revised: 65 mm Hg, serum lactate level > 2 mmol/L, urine output ≥ 5 ml/kg/hr, venous oxygen saturation (SvO2) ≥ 70% and central venous pressure (CVP) of 8 − 12 mmHg.

Original: {+We fill missing values like lab measurements using the last measured value; we clamp real-valued features in between their 0 . 
+}
Revised: {+We fill missing values like lab measurements using the last measured value; we clamp real-valued features in between their 0 . 
+}

Revised: -quantile

Revised: and 0 .

Revised: 95 -quantile values respectively and normalize the features using min-max normalization.

Revised: For training sepsis prediction models, we take the full training cohort but discard the data after the first sepsis onset in sepsis patients, then we label the data per hour, and label the current sepsis outcome as 1 if the true sepsis is going to happen in the next 12 hours (designed for early prediction on sepsis).

Revised: For training shock prediction models, we take the sepsis sub training cohort and discard the data before sepsis onset.

Revised: We also discard the data after septic shock onset in shock patients.

Revised: Then we label in the same way as for sepsis, i.e. label the current shock outcome as 1 if the true shock will take place in the next 12 hrs.

Revised: For those non sepsis patients, we discard the first 12 hrs data after ICU

Revised: [Table]

Revised: [Table]

Revised: admission to reduce data noises and randomly sample a sequence length between 12 hrs up to 7 days per each non sepsis patients.

Revised: More details of data prepossessing are provided in the attached code.

Revised: 5.3 Model Zoo

Revised: Computational cost was measured in ms as the total running time of feeding all the test data (with batch size of 256) calling each individual models on a single GeForce RTX 2080Ti divided by the total number of calls.

Revised: Then we multiply the cost by 10 as the GPU is approximately 10X hardware cost comparing to a CPU.

Revised: Future work can extend the model zoo to include CPU models or running all the models on CPUs based on resource specifications.

Revised: Table 5 and Table 6 respectively show the model prediction AUC scores on the validation set for the sepsis and septic shock stages.

Revised: In addition, we also fit small DKD surrogate models for distilling the predictive probabilities and confidence of the models in the zoo.

Revised: The DKD model is a 4-layer MLP taking the embedding vectors from the first model in each stage, so it obtains similar AUC scores on the validation set comparing to the early models in the zoo but much lower scores comparing to the later heavier models.

Revised: But the mean absolute errors (MAE) of the DKD model on estimating confidence measures of the original models are consistently small, which is beneficial for our soft-gating algorithm that requires only confidence estimation instead of predictive probabilities.

Revised: P r ob a b ilit y o f S e p s i s

Revised: [Figure]

Revised: IDK thresholds

Revised: ICK thresholds

Revised: Predicted Onset

Revised: Real Onset

Revised: P r ob a b ilit y o f

Revised: S e p ti c

Revised: S ho c k

Revised: [Figure]

Revised: We analyze model utilization frequency (the proportion of how many times a model was invoked) in our test cohort and compare model frequencies for hard and soft gating in Figure 4 (models are grouped into 8 groups for Stage 1 and 10 groups for Stage 2): soft gating can skip invocations of many models and directly select the more confident models for faster transitioning to shock stage.

Revised: where ψ is the digamma function defined as the logarithmic derivative of the gamma function, α 0 and α 1 are the concentration parameters estimated by the DKD models.

Revised: More definitions are in Malinin and Gales (2018).

Revised: We show “Entropy of expected” exhibits the best AUC-Cost trade-off path in

Revised: Figure 3.

Revised: 5.5 Qualitative Evaluation

Revised: We walk through an example shock patient’s length of stay in ICU from the test set, and deploy the proposed multi-stage prediction pipeline on it.

Revised: UnfoldML starts the prediction of sepsis with a cheaper model as seen in Figure 5.

Revised: At t=2, the model’s prediction probability reaches the IDK threshold which signifies model’s uncertainty in sepsis prediction.

Revised: Hence, the UnfoldML switches to a costlier and more accurate model (a similar trend is observed at t=4,7).

Revised: At t=5, UnfoldML predicts sepsis onset as the probability of sepsis prediction reaches ICK threshold.

Revised: Note, once sepsis onset is predicted by the cascade, it switches to a cheaper model which predicts septic shock (Stage-2).

Revised: Due to early switching, UnfoldML can detect septic shock significantly earlier.

Revised: In Stage-2, UnfoldML transitions to costlier model once the cheaper model becomes uncertain.

Revised: Lastly, it predicts septic shock once the probability of septic shock detection reaches the ICK 1 threshold.

Revised: Additionally, we randomly slice 35 k time-steps from the sequential data in the test set and visualize them in a TSNE Van der Maaten and Hinton (2008) plot in Figure 6 based on their embedding vectors