Original: UnfoldML : Cost-Aware and Uncertainty-Based Dynamic 2D Prediction for Multi-Stage Classification
Revised: UnfoldML : Cost-Aware and Uncertainty-Based Dynamic 2D Prediction for Multi-Stage Classiﬁcation

Original: Abstract
Revised: Abstract

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

Original: ML models, however, are increasingly more complex, resource intensive, 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 issues are exacerbated for prediction tasks with sequential classification on progressively transitioned stages with “happens-before” relation between them.
Revised: These issues are exacerbated for prediction tasks with sequential classiﬁcation on progressively transitioned stages with “happens-before” relation between them.

Original: 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.
Revised: We argue that it is possible to “unfold” a monolithic single multi-class classiﬁer, typically trained for all stages using all data, into a series of single-stage classiﬁers.

Original: 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.
Revised: Each singlestage classiﬁer can be cascaded gradually from cheaper to more expensive binary classiﬁers that are trained using only the necessary data modalities or features required for that stage.

Original: 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.
Revised: UnfoldML is a cost-aware and uncertainty-based dynamic 2D prediction pipeline for multi-stage classiﬁcation 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: UnfoldML achieves 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 from the 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 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.
Revised: We also show that UnfoldML generalizes to image classiﬁcation, where it can predict different level of labels (from coarse to ﬁne) given different level of abstractions of a image, saving close to 5X cost with as little as 0.4% accuracy reduction.

Original:  Introduction
Revised:  Introduction

Original: 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
Revised: 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

Original: (Feng et al., 2020; Xu et al., 2018), 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: 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.
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 offloaded 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 ofﬂoaded 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 classification models.
Revised: ML classiﬁcation models.

Original: 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.
Revised: 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 “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 identified by a series of stage transitions
Revised: For instance in clinical settings, disease progression is often identiﬁed by a series of stage transitions

Original: Abstract Input
Revised: Abstract Input

Original: Detailed Input
Revised: Detailed Input

Revised: [Figure]

Revised: C

Revised: [Figure]

Revised: s

Original: [Figure]
Revised: [Figure]

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

Original: o c o s tli e
Revised: o c o s tli e

Original: r m od e l s w h e n l ac k
Revised: r m od e l s w h e n l ac k

Original: 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: StageStage-

Original: Stage
Revised: Stage

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: (Sperling et al., 2011; Singer et al., 2016).
Revised: (Sperling et al., 2011; Singer et al., 2016).

Original: Detection on early stage of the disease allows doctors to take 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 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).
Revised: In image classiﬁcations, recognition from super classes to sub classes in a coarse-to-ﬁne manner has shown improved classiﬁcation performance (Dutt et al., 2017; Lei et al., 2017).

Original: 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.
Revised: In all these applications, general multi-class classiﬁers (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 signiﬁcantly high spatio-temporal cost.

Original: 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.
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 ﬁnd the optimal trade-off.

Original: We propose UnfoldML 2 : a cost-aware and uncertainty-based prediction pipeline for dynamic multistage classification.
Revised: We propose UnfoldML 2 : a cost-aware and uncertainty-based prediction pipeline for dynamic multistage classiﬁcation.

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

Original: 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.
Revised: Each single-stage classiﬁer is then cascaded gradually from cheaper to more expensive binary classiﬁers, further reducing the cost by dynamically selecting an appropriate classiﬁer for an input query.

Original: 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.
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-speciﬁed budget limit.

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

Original: 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).
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 conﬁ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).

Original: 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.
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 enables early 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 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.
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 2D query 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 optimal thresholds on the gating functions defined 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 deﬁned for ICK 0 , IDK and ICK 1 .

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

Original: 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.
Revised: To ﬁnd the optimal thresholds, it performs bottom-up grid search over a topologically-sorted list of all the models and identiﬁes the thresholds to minimize the prediction loss, while following a cost constraint.

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

Original: 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.
Revised: It follows a Mixture-of-Expert (MoE) framework to adaptively determine a models’ conﬁdence threshold given all the other model’s conﬁdence in predicting the same query.

Original: 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.
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 classiﬁer 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 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.
Revised: • We design a novel 2D query propagation pipeline that “unfolds” multi-stage prediction workﬂows 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 sufficiently navigate the cost/accuracy tradeoff space and search an optimal set of policies for the designed 2D query propagation.
Revised: • We propose two learning algorithms to sufﬁciently 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 the spatio-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; 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).
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 selection assuming 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 classifiers for optimizing the trade-off between the expected classification error and the feature cost.
Revised: They train a cascade of classiﬁers for optimizing the trade-off between the expected classiﬁcation error and the feature cost.

Original: Early solution (Raykar et 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) applies boosting 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 incorporate deep 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-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.
Revised: In contrast to all of the above work that are only 1-D pipelines for one-step prediction task (can be multi-class classiﬁcations), our method extends to a 2D pipeline that can dynamically forward examples to next steps after they are conﬁdently predicted as passed on the current step.

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

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

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: (1991), for partitioning the training data 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 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).
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 architecture searching (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 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).
Revised: Sparse gates are introduced in MoE so that it can efﬁciently 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 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
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: Confidence criterion has been incorporated into active learning by Li and Sethi (2006) and then extended by Zhu et al.
Revised: Conﬁdence criterion has been incorporated into active learning by Li and Sethi (2006) and then extended by Zhu et al.

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

Original: 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).
Revised: Conﬁdence 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 method posits thresholds on prediction confidence for activating the gates in pipeline expansion.
Revised: Our method posits thresholds on prediction conﬁdence 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 uncertainty distribution in model predictions, which is more computationally efficient than traditional Bayesian approaches (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 efﬁcient 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 so that we only need to run a low-cost multi-head model for producing the weights in MoE efficiently.
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 efﬁciently.

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

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

Original: 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.
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 ∈ 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 .
Revised: Given x 2 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 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 }} .
Revised: If there are a total number of 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 }} .

Original: 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 ) .
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 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
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 conﬁdently know YES ” ( ICK 1 ) that we have correctly identiﬁed the sample at the current stage, otherwise (4) exit the pipeline for those “ I conﬁdently know

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

Original: Figure 1 (a) demonstrates the proposed 2D 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 for each of the three classes of gating functions with the objective of maximizing high system-wide accuracy 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 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.
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 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 .
Revised: These decisions rules can be realized by two groups of parameters: a conﬁdence criterion q : X ⇥ M ! [0 , 1] that measures the conﬁdence 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 conﬁdence score per each prediction made by model m sk .

Original: 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.)
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 within the user-defined 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-deﬁned cost budget, or an ICK 1 class s.t.

Original: the system will transition the query to the next stage 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 the decision 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 confidence of model m sk ’s prediction on data x t .
Revised:  where q sk ( x t ) is a short notation for q ( x t , m sk ) measuring the conﬁdence of model m sk ’s prediction on data x t .

Original: The goal of configuring an optimal pipeline given a restricted computation resource can be formalized as the following optimization problem:
Revised: The goal of conﬁguring 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 prediction loss on data D and c is a user-specified 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-speciﬁed 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 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 =
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 = {  x i , ( y 1 i , t 1 i ) , · · · , ( y Si , t Si )  } Ni =1 , where x i =

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

Original: 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:
Revised: We ﬁrst 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-allocated for stage- s satisfying(cid:80) s
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

Original: c s = c .
Revised: [Equation]

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

Original: (1) into two 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 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.
Revised: We can efﬁciently learn each stage’s optimal IDK gate parameters by solving Sub-Objective 1, and learn the optimal ICK 1 gate parameters by ﬁxing 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 confidence in the prediction is below a threshold.
Revised: the gate is only activated if the level of conﬁdence in the prediction is below a threshold.

Original: [Equation]
Revised: [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 s can 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 as
Revised: Now we write the conditional probability of being at stage- s given input x t as

Revised: Pr ( y s = 1 | x

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where m sk ( x t ) =
Revised: where m sk ( x t ) =

Original: Pr (cid:0) y s = 1 | x t ; m sk (cid:1) .
Revised: Pr  y s = 1 | x t ; m sk  .

Original: The one-stage loss function in Hard-Gating is defined as the negative log-likelihood loss
Revised: The one-stage loss function in Hard-Gating is deﬁned as the negative log-likelihood loss

Original: [Equation]
Revised: [Equation]

Original: where p sit =
Revised: where p sit =

Original: Pr ( y s = 1 | x it ; m casc s ) , y sit =
Revised: Pr ( y s = 1 | x it ; m casc s ) , y sit =

Original: 1 only if y si = 1 and t ∈ [ t 1 i − δ t , t 1 i ] for some δ t time-steps we wish to early detect the next stage s
Revised: 1 only if 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

Original: + 1 .
Revised: + 1 .

Original: Algorithm 1 in Appendix 5.1 describes a bottom-up grid search algorithm for for learning hard-gating parameters.
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 order 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 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
Revised: For any stage s , Algorithm 1 starts by assigning all the N s samples for that stage to the ﬁrst 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’s model in the list until the cost exceeds the user-defined 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-deﬁned budget c s for that stage.

Original: In each iteration of searching the cutoff α sk , we set an upper bound maxA on the maximum searching value to avoid over-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 overfit and put all the samples into the IDK class, then assign them to the next level’s model.
Revised: Without setting this bound, the algorithm could overﬁt 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 an objective function that can be efficiently 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 efﬁciently solved using gradient descent algorithms.

Original: 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.
Revised: 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.

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 we define Soft-gating as
Revised: Formally we deﬁne Soft-gating as

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 defined as a mixture of the K s models available for stage- s prediction:
Revised: s given input x t is deﬁned as a mixture of the K s models available for stage- s prediction:

Original: [Equation]
Revised: [Equation]

Original:  t )) .
Revised: ( q sj ( x t )) .

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 models in the zoo, which conflicts 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 conﬂicts 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 the prediction confidence for each model in the zoo.
Revised: Knowledge Distillation (DKD) for training a small surrogate model for each stage to quantify the prediction conﬁdence 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 smaller distilled model only needs to be run once per query, requiring much less cost than running all the models.
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 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
Revised: In our experiment, we utilize the ﬁrst 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 posit a Dirichlet prior distribution over the parameters π characterizing the predicted output categorical distribution (i.e., binomial in our setup) and a surrogate prior network f is fit to generate the concentration parameters α sk in the prior:
Revised: The idea of DKD is to posit a Dirichlet prior distribution over the parameters ⇡ characterizing the predicted output categorical distribution (i.e., binomial in our setup) and a surrogate prior network f is ﬁt to generate the concentration parameters ↵ sk in the prior:

Original: [Equation]
Revised: [Equation]

Original: If the learnt concentration parameters yield a flat prior distribution, 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 ﬂat prior distribution, 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, [-
-] + α sk, 1 ) .
Revised: ↵ sk, 1 / ( ↵ sk, [-
-]0 + ↵ sk, 1 ) .

Original: For training the DKD model, which 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 define the loss function for head k as a Kullback-Leibler (KL) divergence between the prior distribution and empirical observed distribution:
Revised: MLP per each stage in this paper, we deﬁne the loss function for head k as a Kullback-Leibler (KL) divergence between the prior distribution and empirical observed distribution:

Original: [Equation]
Revised: [Equation]

Original: where p sk ( x it ) are the real predicted probabilities produced from model
Revised: where p sk ( x it ) are the real predicted probabilities produced from model

Original: m sk on input x it .
Revised: m sk on input x it .

Original: Additionally, we also add the cross-entropy loss 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 selected models do not exceed user-imposed cost constraint by enforcing sparse gating weights over the model choices.
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
Revised: Therefore, we reformulate the objective function in

Original: Hard-Gating Sub-Objective 1 as the following Lagrangian function
Revised: Hard-Gating Sub-Objective 1 as the following Lagrangian function

Original: [Equation]
Revised: [Equation]

Original: where the second term L s cost takes the cost constraint, controlled by λ > 0 and the third term
Revised: where the second term L s cost takes the cost constraint, controlled by  > 0 and the third term

Original: L s sparse imposes further sparse regularization on the gating weights w.r.t their L 1 norms, controlled by µ > 0 .
Revised: L s sparse imposes further sparse regularization on the gating weights w.r.t their L 1 norms, controlled by µ > 0 .

Original: 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
Revised: Formally, we write the last two terms as L s cost =  max(0 , cost ( m casc s )  c s )  2 and

Original: L s sparse =
Revised: L s sparse =

Original: (cid:80) i,t ∥ G IDK (ˆ q sk ( x it ))
Revised: P i,t k G IDK (ˆ q sk ( x it ))

Original: ∥ 1 .
Revised: k 1 .

Original: Now we use stochastic gradient descent algorithms to learn the optimal a s and b s minimizing the above 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-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.
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 ICKare 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 finds 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 ﬁnds the optimal threshold achieving the most left upper corner in the ROC curve.

Original: Finally, Algorithm 3 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 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).
Revised: The ﬁrst 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 filters out queries that are not of interest in Stage-1 and then refines the predictions into fine classes in Stage-2.
Revised: 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ﬁnes the predictions into ﬁne 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 deidentified health records from over 50 , 000 critically ill patients who stayed in the ICUs of the Beth Israel
Revised: The database consists of deidentiﬁed 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.

Original: The final cohort includes a total of 34 , 475 ICU patients, from which 2 , 370 (6 . 8%) presented with
Revised: The ﬁnal 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 cohort of 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 the training 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 train the 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 classifiers 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 classiﬁers for Stage-1 (Sepsis) and for Stage-2 (Septic Shock) using a variety of architectures and features.

Original: 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).
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 hidden size 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 patient data 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 with different 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 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.
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.

Original: Confidence Measures We consider four choices for measuring the confidence q of a model prediction.
Revised: Conﬁdence Measures We consider four choices for measuring the conﬁdence q of a model prediction.

Original: Given p , the model’s predictive probability of output Y being 1 , we define
Revised: Given p , the model’s predictive probability of output Y being 1 , we deﬁne

Original: • 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: [-− α 
-]
Revised: • Max probability: max  1  p, p  , • Entropy:  p · log p + (1  p ) · log(1  p )  , • Entropy of expected: [-− α 
-] ↵ 0 ↵ 0 + ↵ 1 ·  ( ↵ 0 )

Original: [-α 
 + α 1 · (cid:0) ψ-] ( α 0 ) 
Revised: [-α 
 + α 1 · (cid:0) ψ-] ( ↵ 1 )

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

Original: Spatio-Temporal Cost.
Revised: Spatio-Temporal Cost.

Original: Given the trained model zoo, we profile spatio-temporal costs for each of the models.
Revised: Given the trained model zoo, we proﬁle 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 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 cloud offerings 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 first system to provide 2-
Revised: To our knowledge, UnfoldML is the ﬁrst system to provide 2-dimensional cascading

Original: dimensional cascading predictions.
Revised: predictions.

Original: 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.
Revised: A reasonable baseline for our method are models that frame the multi-stage sequential task as a multi-class classiﬁcation 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 clinical time-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 changing the 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 unified multi-class model works end to end for predicting the multi classes.
Revised: One uniﬁed multi-class model works end to end for 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 similar hard-gating algorithm to learn hard thresholds on prediction entropy for making decisions if the system 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 its horizontal 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 them on the two single-stage classification tasks on sepsis and septic shock predictions respectively.
Revised: We evaluate them on the two single-stage classiﬁcation tasks on sepsis and septic shock predictions respectively.

Original: • 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.
Revised: • 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 is well-known in early detection of septic shock; GPU-based models such as LSTMs.

Original: • IDK-cascade (Wang et al., 2017):
Revised: • IDK-cascade (Wang et al., 2017):

Original: Prior work of 1-D prediction cascade on IDK classes using a hard-gating like algorithm.
Revised: 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 Classification Performance.
Revised: End-to-End Classiﬁcation Performance.

Original: 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.
Revised: To evaluate the prediction performance on our multi-stage task, we treat UnfoldML as a multi-class classiﬁer which predicts one of the same three classes as deﬁned in Baseline.

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

Original: Each run contributes a point in 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 and
Revised: We compute the convex hull of the searched points in the set and

Original: [Figure]
Revised: [Figure]

Revised: nd - t

Revised: [Figure]

Original: E nd
Revised: E nd

Original: AU C
Revised: AU C

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure][Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure][Figure]
Revised: [Figure] [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: demonstrate the resulting area of the hull in Figure 2a.
Revised: demonstrate the resulting area of the hull in Figure 2a.

Original: We train our baseline method repeatedly by varying 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 significantly higher AUC regions than the multi-class LSTM baseline and performs more consistently than the hard-gating algorithm.
Revised: In comparison, UnfoldML -S ( soft-gating ) searches signiﬁcantly 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).

Original: UnfoldML -S presents the highest score comparing to the baseline with UnfoldML
Revised: 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 2.
Revised: Table 2.

Original: 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.
Revised: The recommended conﬁguration 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 configuration for our pipeline
Revised: As shown in Table 2, an optimal conﬁguration for our pipeline

Original: 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.
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 each stage, 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 AUC gain 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 UnfoldML is a result of the dynamic nature of inference pipeline, which selects the optimal pathway for each patient.
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 when they are confident enough, in contrast, baselines ( UnfoldML -A, Multi-Class LSTM-A) always use the most accurate model.
Revised: The savings in spatio-temporal cost can be attributed to the usage of low cost models when they are conﬁdent enough, in contrast, baselines ( UnfoldML -A, Multi-Class LSTM-A) always use the most accurate model.

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 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.
Revised: For both predictions, UnfoldML achieves a signiﬁcant 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 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.
Revised: Mutli-class LSTM-A. It beneﬁts 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 Classification
Revised: 4.2 Task 2: Subcategory Classiﬁcation

Original: To demonstrate the generalizability of UnfoldML we conduct experiments on a computer vision task where we wish to detect a subcategory-of-interest .
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: In the subcategory-of-interest task, we wish to
Revised: In the subcategory-of-interest task, we wish to

Original: [Table] [Table]
Revised: [Table] [Table]

Original: 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.
Revised: accurately predict if an image falls within a speciﬁc subcategory of a dataset with classes that form a hierarchical structure, while still ensuring our spatio-temporal costs remain within the user-deﬁned budget.

Original: 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.
Revised: 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 using the
Revised: For our experiment using the

Original: 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.
Revised: CIFAR-100 dataset (Krizhevsky et al., 2009), we deﬁne two separate subcategory-of-interest tasks based on the real-world system applications of computer vision systems.

Original: 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).
Revised: We want to identify images of a speciﬁc subcategory from one of two chosen categories: ‘people’ (baby, boy, girl, man, woman) and ‘vehicles’ (bicycle, bus, motorcycle, pickup truck, train).

Original: CIFAR-100 consists of 60,000 32x colour images in total of 100 ‘fine’ classes.
Revised: CIFAR-100 consists of 60,000 32x colour images in total of 100 ‘ﬁne’ classes.

Original: There are 500 training images and 100 testing images per class.
Revised: There are 500 training images and 100 testing images per class.

Original: In addition, each image also is assigned a ‘coarse’ label indicating which category (such as people, vehicles, trees, etc) it belongs to.
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 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 .
Revised: In our experiment, we do a multi-class classiﬁcation on the ’coarse-granularity’ labels in Stage-1 to detect the category, then we do multi-class classiﬁcation on the original ‘ﬁne-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 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: Stage-1 we train K 1 number of multi-class classiﬁcation models using the full dataset of CIFARwith coarse-granularity labels as our target subcategories, resulting in a total of 20 classes.

Original: We train models for Stage-1 using the full 32x32 image and as well as random crops of size 24x24 or 16x16.
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-2 we train K 2 multi-class models to identify the next 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 our Stage-2 models only using data with the coarsegranularity label for our subcategory in Stage-1.
Revised: We train our Stage-2 models only using data with the coarsegranularity label for our subcategory in Stage-1.

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

Original: We profile each model by computing the number of Multiply-Accumulate
Revised: We proﬁle each model by computing the number of Multiply-Accumulate

Original: Operations (MACs) and use it as proxy measure of our spatio-temporal cost.
Revised: Operations (MACs) and use it as proxy measure of our spatio-temporal cost.

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

Original: We compare UnfoldML against a multi-class WRN classifier baseline which directly predicts labels at the fine-granularity.
Revised: We compare UnfoldML against a multi-class WRN classiﬁer baseline which directly predicts labels at the ﬁne-granularity.

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

Original: We evaluate both accuracy and MACs, and report the results in Table 4.
Revised: We evaluate both accuracy and MACs, and report the results in Table 4.

Original: The results show that UnfoldML is able to find the optimal gating policy for the pipeline and achieve
Revised: The results show that UnfoldML is able to ﬁnd the optimal gating policy for the pipeline and achieve

Original: [Table]
Revised: [Table]

Original: 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% reduction in accuracy for ‘vehicle’ subcategory-of-interest identification.
Revised: a spatio-temporal cost-savings of 6.9X with 1% reduction in accuracy for ‘people’ subcategory classiﬁcation; we achieve a cost-savings of 4.7X with 0.4% reduction in accuracy for ‘vehicle’ subcategory-of-interest identiﬁcation.

Original: Thus, UnfoldML achieves a savings in spatio-temporal cost in for this task without compromising the overall accuracy.
Revised: Thus, UnfoldML achieves a savings in spatio-temporal cost in for this task without compromising the overall accuracy.

Original: UnfoldML can accomplish this because when it is confident that queries do not belong to our subcategory-of-interest ( ICK 0 ), we are able to early exit the query from the prediction pipeline.
Revised: UnfoldML can accomplish this because when it is conﬁdent that queries do not belong to our subcategory-of-interest ( ICK 0 ), we are able to early exit the query from the prediction pipeline.

Original: When UnfoldML is not confident ( IDK ), it upgrades queries to costlier and more accurate models.
Revised: When UnfoldML is not conﬁdent ( IDK ), it upgrades queries to costlier and more accurate models.

Original: Only when UnfoldML can confidently predict the subcategory ( ICK 1 ) in Stage-1 will it transition the query to Stage-2 and make the final prediction.
Revised: Only when UnfoldML can conﬁdently predict the subcategory ( ICK 1 ) in Stage-1 will it transition the query to Stage-2 and make the ﬁnal prediction.

Original: This results in a significant cost-savings without compromising accuracy.
Revised: This results in a signiﬁcant cost-savings without compromising accuracy.

Original: A recent alternative method for trading-off accuracy and cost in image classification is CwCF (Classification with Costly Features)
Revised: A recent alternative method for trading-off accuracy and cost in image classiﬁcation is CwCF (Classiﬁcation with Costly Features)

Original: (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: (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 classiﬁer.

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

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

Original: In order to establish a fair baseline against UnfoldML , we used the most accurate multi-class WRN for the HPC 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 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.
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 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.
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:  Conclusion
Revised:  Conclusion

Original: ML models, including for healthcare applications, are growing exponentially in size and cost of inference.
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, monolithic models require expensive hardware, and can not fit on bed-side compute or even on site compute clusters with clinical implications.
Revised: Costlier, monolithic models require expensive hardware, and can not ﬁt on bed-side compute or even on site compute clusters with clinical implications.

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

Original: 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.
Revised: It implements a query propagation mechanism 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.

Original: 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
Revised: Each query is allowed to (1) conﬁdently 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 thresholds at 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 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.
Revised: UnfoldML builds on a fundamental insight that classes may have a “happens before” relationship between them, which can be leveraged to “unfold” a classiﬁer, leading to savings in spatio-temporal cost (how much resource used for how long) and clinically signiﬁcant earlier onset prediction.

Original: UnfoldML improves the frontier of optimality in 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 performance for 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 mere 0.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.

Original: Acknowledgements
Revised: Acknowledgements

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

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

Original: We would also like to acknowledge Dr.
Revised: We would also like to acknowledge Dr.

Original: 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: 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.

Original: Disclaimer:
Revised: Disclaimer:

Original: 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.Appendix 5.1 Training Algorithms
Revised: Any opinions, ﬁndings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reﬂect the views of the National Science Foundation.

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

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

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

Original: Input

Original: Output

Original: 2: α ∗ s =

Original: ModelAssign cost

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

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

Original: Idx 4 k ← ∪ I ( ModelAssign t

Original: 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) ) .

Original: 11: α ∗ sk ← minQ .

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

Original: IDK ← ∪ Idx 4 k

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

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

Original: [ m k ] >

Original: c s then 16:

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

Original: 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

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

Original: 5.2 Data preparation

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

Original: (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.

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

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

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

Original: 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.

Original: 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 .

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

Original: Then according to Singer et al.

Original: , we identify the onset of septic shock as

Original: Output

Original: 1: procedure S OFT G ATING

Original: ( 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

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

Original: (cid:1)

Original: ▷ distillation 5:

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

Original: ||6:

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

Original: s sparse 7:

Original: Optimize L a s , using

Original: Reduce lr by factor 0 .

Original: 5 once learning stagnates.

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

Original: Input

Original: Output

Original: RAINING ( D

Original: M ) 2:

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

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

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

Original: ← HardGating( D s

Original: M s c s

Original: C q ) 6: or, a

Original: , b

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

Original: Step 2: Learn ICK 1 gate parameters.

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

Original: [Equation]

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

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: 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).

Original: We result at 229(9 .

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

Original: 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.

Original: 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.

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

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

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

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

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

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

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

Original: 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 . 
-]

Original: -quantile

Original: and 0 .

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

Original: 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).

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

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

Original: 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.

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

Original: [Table]

Original: [Table]

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

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

Original: 5.3 Model Zoo

Original: 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.

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

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

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

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

Original: 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.

Original: 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.

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

Original: [Figure]

Original: IDK thresholds

Original: ICK thresholds

Original: Predicted Onset

Original: Real Onset

Original: P r ob a b ilit y o f

Original: S e p ti c

Original: S ho c k

Original: [Figure]

Original: 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.

Original: 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.

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

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

Original: Figure 3.

Original: 5.5 Qualitative Evaluation

Original: 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.

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

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

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

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

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

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

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

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

Original: 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