Original: UnfoldML : A Cost-Aware 2-D Dynamic Prediction Pipeline for Multi-Stage Classiﬁcation
Revised: UnfoldML : A Cost-Aware 2-D Dynamic Prediction Pipeline for Multi-Stage Classiﬁcation

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AfﬁliationAddress email
Revised: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

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

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

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

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

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

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

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

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

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

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

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

Original:  Introduction
Revised:  Introduction

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

Original: pedestrian detection
Revised: pedestrian detection

Original: Zhang et al., 2018;
Revised: Zhang et al., 2018;

Original: Cai et al., 2015), etc.
Revised: Cai et al., 2015), etc.

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

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

Original: Run time environment like mobile devices or bedside-patient monitors are commonly resource constrained, and applications that can be ofﬂoaded to cloud computing always aim for a reduced cloud bill.
Revised: Run time environment like mobile devices or bedside-patient monitors are commonly resource constrained, and applications that can be 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 classiﬁcation models
Revised: ML classiﬁcation models

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

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

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

Original: Abstract Input
Revised: Abstract Input

Original: Detailed 
Revised: Detailed Input

Original: [Figure]

Original: IDK IDK
Revised: IDK IDK

Original: ICK
Revised: ICK

Original: [Figure]
Revised: [Figure]

Revised: Exit

Revised: [Figure]

Original: C a s
Revised: C a s

Original: [Figure]
Revised: [Figure]

Original: 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 r m od e l s w h e n l ack
Revised: o c o s t l i e r m od e l s w h e n l ac

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

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

Original: [Figure]
Revised: [Figure]

Original: ICK
Revised: ICK

Revised: [Figure]

Original: ICK
Revised: ICK

Original: [Figure]
Revised: [Figure]

Original: ICK ICK
Revised: ICK ICK

Original: [Figure]
Revised: [Figure]

Revised: ICK

Revised: ICK

Revised: ICK

Revised: [Figure]

Revised: [Figure]

Original: appropriate actions in time before it enters into late severe stages.
Revised: appropriate actions in time before it enters into late severe stages.

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

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

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

Original: We propose UnfoldML : a cost-aware prediction pipeline that leverages the “happens-before” propertyin multi-stage classiﬁcation.
Revised: We propose UnfoldML : a cost-aware prediction pipeline that leverages the “happens-before” propertyin multi-stage classiﬁcation.

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

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

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

Original: It computes a classiﬁer’s prediction conﬁdenceon a query then directs the query through one the following three gates: 1) “ I conﬁdently know
Revised: It computes a classiﬁer’s prediction conﬁdenceon 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’tknow ” ( IDK ), which upgrades the query to a higher accuracy but costlier model, producing a moreconﬁdent result within the budget (vertical cascading); and 3) “ I conﬁdently know YES ” ( ICK 1 ),which transitions the query to the next stage of the prediction task (horizontal forwarding).
Revised: NO ” ( ICK 0 ), which rejects the current query and early exits from the pipeline (exit); 2) “ I don’tknow ” ( IDK ), which upgrades the query to a higher accuracy but costlier model, producing a moreconﬁdent result within the budget (vertical cascading); and 3) “ I conﬁdently know YES ” ( ICK 1 ),which transitions the query to the next stage of the prediction task (horizontal forwarding).

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

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

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

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

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

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

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

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

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

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

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

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

Original: prediction pipeline with minimal accuracy degradation.•
Revised: prediction pipeline with minimal accuracy degradation.

Original:  We propose two learning algorithms to sufﬁciently navigate the cost/accuracy tradeoff space andsearch an optimal set of policies for the designed 2-D 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 2-D query propagation.•

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

Original:  Related Work
Revised:  Related Work

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

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

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

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

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

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

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

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

Original: Further, we also develop a more efﬁcient 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 decentalgorithms for learning the parameters efﬁciently.
Revised: Mixture-of-Experts (MoE) modeling and knowledge distillation, which can apply gradient decentalgorithms 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: ([n. d.]), for partitioning the trainingdata and feeding them into separate neural networks during the learning process.
Revised: ([n. d.]), for partitioning the trainingdata and feeding them into separate neural networks during the learning process.

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

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

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

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

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

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

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

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

Original: Our methodposits thresholds on prediction conﬁdence for activating the gates in pipeline expansion.
Revised: Our methodposits 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 uncertaintydistribution in model predictions, which is more computationally efﬁcient than traditional Bayesianapproaches (MacKay, 1992; Mackay, 1992; Hinton and Van Camp, 1993).
Revised: Prior Networks (BPNs) (Malinin and Gales, 2018) have been proposed to estimate the uncertaintydistribution in model predictions, which is more computationally efﬁcient than traditional Bayesianapproaches (MacKay, 1992; Mackay, 1992; Hinton and Van Camp, 1993).

Original: We propose Dirichlet
Revised: We propose Dirichlet

Original: Knowledge Distillation (DKD) based on BPNs for distilling prediction uncertainty in large models sothat we only need to run a low-cost multi-head model for producing the weights in MoE efﬁciently.
Revised: Knowledge Distillation (DKD) based on BPNs for distilling prediction uncertainty in large models sothat 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 2-D prediction pipeline UnfoldML , which learns the optimal policy formaking “ I conﬁdently know ” ( ICK ) predictions on sequential multi-stage classiﬁcation tasks.
Revised: We introduce a dynamic 2-D prediction pipeline UnfoldML , which learns the optimal policy formaking “ I conﬁdently know ” ( ICK ) predictions on sequential multi-stage classiﬁcation tasks.

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

Original: 3.1 Problem Formulation
Revised: 3.1 Problem Formulation

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

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

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

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

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

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

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 144 example x t coming at time t and the current model choice m sk as an input to determine whether themodel for the query should (1) upgrade ‘vertically’ to costlier models in the same stage, (2) transition‘horizontally’ to the next stage, or(3) return as rejected and exit the pipeline.
Revised: X ⇥ M ! M , which takesexample x t coming at time t and the current model choice m sk as an input to determine whetherthe model for the query should take one of the aforementioned three actions: upgrade vertically,transition horizontally, or exit the pipeline.

Original: These decisions rulescan be realized by two groups of parameters: a conﬁdence criterion q : X ⇥ M !
Revised: These decisions rules can be realized by two groupsof parameters: a conﬁdence criterion q : X ⇥ M !

Original: [0 , 1] that148 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 predictionmade by model m sk .
Revised: [0 , 1] that measures the conﬁdence score of amodel’s prediction on a data example, and two gate functions G IDK , G ICK 1 : [0 , 1] ⇥ { True , False }that are applied on to the conﬁdence score per each prediction made by model m sk .

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

Original: the system will upgrade the queryto a costlier model while remaining within the user-deﬁned cost budget, or an ICK 1 class s.t. thesystem will transition the query to the next stage in the pipeline.
Revised: the system will upgrade the query to a costlier model while remaining withinthe user-deﬁned cost budget, or an ICK 1 class s.t. the system will transition the query to the nextstage in the pipeline.

Original: The third gate G ICK 0 is determinedby ¬ 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 thedecision rule used at each prediction stage as follows

Original: [Equation]
Revised: [Equation]

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

Original: The goal of conﬁguring an optimal pipeline given a restricted computation resource canbe formalized as the following optimization problem, which minimizes an end-to-end prediction loss
Revised: The goal of conﬁguring an optimal pipeline given a restricted computation resource canbe formalized as the following optimization problem:

Original: L of m casc on an given observed data set D with respect to (w.r.t) the gate speciﬁcations G s.t. thetotal cost of running the pipeline remains within the user-speciﬁed cost-constraints for the systems.t.
Revised: [Equation]

Original: min G L  m casc ; D ) s.t. cost ( m casc ; D )  c , where G consists of  two gate functions G IDKand G ICK 1 . Next, we introduce two training algorithms forlearning the optimal gate functions in UnfoldML .
Revised: where G consists of the two gate functions G IDK and G ICK 1 , L denotes the end-to-end predictionloss on data D and c is a user-speciﬁed cost-constraint for the system.

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 = {  x i , ( y 1 i , t 1 i ) , · · · , ( y Si , t Si )  } Ni =1 , where(1) x i = ( x i 1 , · · · , x iT i )is an input sequence observed for individual i ,(2) y si 2 { 0 , 1 } indicates whether the example entered to stage- s , and if yes , we use(3) t si 2 ; [ [1 , T i ] to determine when it entered.
Revised: Given a training data set D , which does not include any data that was used in the training of the modelsin our model zoo: D = {  x i , ( y 1 i , t 1 i ) , · · · , ( y Si , t Si )  } Ni =1 , where x i = ( x i 1 , · · · , x iT i ) is an inputsequence observed for individual i , y si 2 { 0 , 1 } indicates whether the example entered to stage- s ,and if yes , we use t s i 2 ; [ [1 , T i ] to determine when it entered.

Original: 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) }
Revised: We ﬁrst partition the multi-stage datainto S one-stage data sets: D s = { ( x i [ t s  1 i : t si ] , y si == 1) }

Original: [ { ( x i [ t s  1 i : T i ] , y si == 0) } , then divide the learning of IDK and ICK gate parameters into twoseparable sub-problems: Sub-Objective 1: min G IDKs L s  m casc s ; D s  s.t. cost ( m casc s ; D s )  c s , s =1 , · · · , S and Sub-Objective 2: min G ICK 1 L  m casc ; D , G IDK ⇤  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
Revised: [ { ( x i [ t s  1 i : T i ] , y si == 0) } , then dividethe learning of IDK and ICK gate parameters into two separable sub-problems: [Equation] where L s is the one-stage prediction loss on data D s , and c s is the cost budget that is pre-allocatedfor stage- s satisfying

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

Original: Decomposing the end-to-end optimization problem in Soft-Gating into two sub-problems HardGating allows us to parallelize the training process and efﬁciently learn each stage’s optimal IDKgate parameters by solving Sub-Objective 1.
Revised: Decomposing the end-to-end optimization problem in Eq. (1) intotwo sub-problems in Eq. (2) allows us to parallelize the training process.

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

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

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

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

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

Original: =
Revised: =

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

Revised: ↵ sk ) ,

Revised: [Equation]

Original: ↵ sk ) , where I ( · ) is an indicator function.
Revised: -gating [Equation] [Equation] 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 scan 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 giveninput x t as Pr ( y s = 1 | x t ; m [Equation]
Revised: Now we write the conditional probability of being at stage- s given input x t as Pr ( y s = 1 | x t ; m casc s ) = P

Original: sk
Revised: K s k

Original: [Equation]{+=
 I ( q sk ( x t ) ↵ sk ) · Q k  1 j =1 I ( q sj ( x t ) <+}
Revised: {+=
 I ( q sk ( x t ) ↵ sk ) · Q k  1 j =1 I ( q sj ( x t ) <+}

Original:  x [Equation] ) ,where m sk ( x t ) =
Revised: ↵ sj ) · m sk ( x t ) , where m sk ( x t ) =

Original: Pr  y s = 1 | x t ; m sk  .
Revised: Pr  y s = 1 | x t ; m sk  .

Original: The one-stage loss function in Hard-Gating is deﬁned asthe negative log-likelihood loss 
Revised: The onestage loss function in Hard-Gating is deﬁned as the negative log-likelihood loss L s nll  m casc s ; D s  =

Original: L
Revised: P

Original: [Equation]
Revised: N s i =

Original: nll
Revised: P T i t =1 y sit · log p sit +(1  y sit ) ·

Original: [Equation] ,where p sit =
Revised: log(1  p sit ) , where p sit =

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

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

Original: We sort the model list M s = { m s 1 , · · · m sK s } in a monotonically increasing order 190 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 185 w.r.t AUC and cost which discards any sub-optimal models.

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

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

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

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

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

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

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

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

Original: In this algorithm,we deﬁne the gate function G IDK to be a ReLU function parameterized by a pair of coefﬁcients( a sk , b sk ) s.t.
Revised: In this algorithm,we 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 )
Revised: the gate is only activated if the linear product a sk · q sk ( x t )

Original:  b sk > 0 .
Revised:  b sk > 0 .

Original: Formallywe deﬁne Soft-gating as G IDK ( q sk ( x t ))
Revised: Formally wedeﬁne Soft-gating as Soft [Equation] gating

Original: = ReLU ( a sk · q sk ( x t )  b sk ) .
Revised: [Equation]

Original: Therefore the conditionalprobability of being at the stage- s given input x t is deﬁned as a mixture of the K s models availablefor stage- s prediction:
Revised: Therefore the conditional probability of being at the stage- s given input x t is deﬁned as a mixture ofthe K s models available for stage- s prediction:

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

Original: [Equation]
Revised: [Equation]

Original: Now, the negative loglikelihood 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 normalizationterm in the mixture of experts requires running all the candidates models in the zoo, which conﬂictswith our cost-saving goal.
Revised: However, the normalization term in the mixture of experts requires running all the candidates modelsin the zoo, which conﬂicts with our cost-saving goal.

Original: Therefore, we propose using a Dirichlet Knowledge Distillation (DKD) fortraining a small surrogate model for each stage to quantify the prediction conﬁdence for each modelin the zoo.
Revised: Therefore, we propose using a Dirichlet Knowledge Distillation (DKD) for training a small surrogate model for each stage to quantify theprediction 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 expertmodels 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 runonce per query, requiring much less cost than running all the models.
Revised: The smallerdistilled model only needs to be run once per query, requiring much less cost than running all themodels.

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

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

Original: [Equation]
Revised: [Equation]

Original: A surrogate priornetwork f is ﬁt to generate the concentration ↵ sk of the Dirichlet distribution.
Revised: surrogate priornetwork f is ﬁt to generate the concentration ↵ sk of the Dirichlet distribution.

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

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

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

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

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

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

Original: Kullback-Leibler (KL) divergence between the prior distribution and empirical observed distribution: L ( ↵
Revised: Kullback-Leibler (KL) divergence between the prior distribution and empirical observed distribution: L ( ↵

Original: [Equation]
Revised: [Equation]

Original:  )= [Equation]  , where p sk ( x it ) are the real predictedprobabilities produced from model m sk on input x it .
Revised: x it )  , where p sk ( x it ) are the real predictedprobabilities produced from model m sk on input x it .

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

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

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

Original: Therefore, we reformulate the objective function in Hard-Gating Sub-Objective 1 as thefollowing Lagrangian function min G IDK L s nll ( m casc ; D s )
Revised: Therefore, we reformulate the objective function in Hard-Gating Sub-Objective 1 as thefollowing Lagrangian function min G IDK L s nll ( m casc ; D s )

Original: +  L s cost + µ L s sparse , where the second term
Revised: +  L s cost + µ L s sparse , where the second term

Original: L s cost takes the cost constraint, controlled by  > 0 and the third term L s sparse imposes further sparseregularization on the gating weights w.r.t their L 1 norms, controlled by µ > 0 .
Revised: L s cost takes the cost constraint, controlled by  > 0 and the third term L s sparse imposes further sparseregularization on the gating weights w.r.t their L 1 norms, controlled by µ > 0 .

Original: Formally, we writethe last two terms as L s cost =[Equation] and L
Revised: Formally, we writethe last two terms as L s cost =

Original: [Equation]
Revised: [Equation]

Original: Now we use stochastic gradient descent algorithms to learn the optimal a s and b s minimizing theabove loss.
Revised: Now we use stochastic gradient descent algorithms to learn the optimal a s and b s minimizing theabove 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: Hard-Gating (3) Sub

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

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

Original: the type I error on class ICK 1 and type IIerror on ICK 0 are both minimized.
Revised: the type I error on class ICK 1 and type II error on ICK 0are 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 errorand type II error in various ways, we pick the Closet-to-(0,1) (Perkins and Schisterman, 2006) methodthat ﬁnds the optimal threshold achieving the most left upper corner in the ROC curve.
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 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 canuse either Algorithm 1 or 2 to learn the optimal gating policy.
Revised: Finally, Algorithm3 in Appendix 5.1 gives the end-to-end training Algorithm for UnfoldML , where we can use either Algorithm 1 or 2 to learn the optimal gating policy.

Original:  Experiments
Revised:  Experiments

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

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

Original: The second task is detection of the subcategoryof-interest in a label hierarchy that ﬁlters out queries that are not of interest in Stage-1 and then reﬁnesthe predictions into ﬁne classes in Stage-2.
Revised: The second task is detection of the subcategoryof-interest in a label hierarchy that ﬁlters out queries that are not of interest in Stage-1 and then reﬁnesthe 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 deidentiﬁedhealth records from over 50 , 000 critically ill patients who stayed in the ICUs of the Beth Israel
Revised: The database consists of deidentiﬁedhealth 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 ﬁnal 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 cohortof patient data into a training set ( 70% ), validation set ( 20% ) and test set ( 10% ).
Revised: We randomly split our cohortof patient data into a training set ( 70% ), validation set ( 20% ) and test set ( 10% ).

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

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

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

Original: 4.1.1 Experimental Setup
Revised: 4.1.1 Experimental Setup

Original: Model Zoo.
Revised: Model Zoo.

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

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

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

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

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

Original: All models start with withbasic demographic features and vital signs, and are then extended with additional features includingresults from lab tests, beside monitoring, and medication/IV treatments.
Revised: All models start with withbasic demographic features and vital signs, and are then extended with additional features includingresults from lab tests, beside monitoring, and medication/IV treatments.

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

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

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 (spatialcost).
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

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

Original: Baseline.
Revised: Baseline.

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

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

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

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

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

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

Original: One uniﬁed multi-class model works end to end 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 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 usesa similar
Revised: It uses

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Original: In comparison, UnfoldML -S( soft-gating ) searches signiﬁcantly higher AUC regionsgiven different cost values 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 andperforms more consistently than the hard-gating algorithm.

Original: To quantify thetrade-off evaluation in Figure 2a, we calculate the area of the convex hull searched by each method(Table 4.1.2).
Revised: To quantify the trade-off evaluation in 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]
Revised: [Table]

Original: Table 2.
Revised: Table 2.

Original: 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 thatform the top-left corner of the hull.
Revised: The recommended conﬁguration for systems that effectively trade off spatio-temporalcost and end-to-end AUC is to select models along the frontier line drawn through the points thatform the top-left corner of the hull.

Original: As shown in Table 2, an optimal conﬁguration 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 accuratebaseline Multi-Class LSTM-A ( 88 . 8% ), while operating at a 19.6X spatio-temporal cost reduction.
Revised: UnfoldML -S (Rec.) can achieve comparable end-to-end AUC score ( 88 . 9% ) with the most accuratebaseline Multi-Class LSTM-A ( 88 . 8% ), while operating at a 19.6X spatio-temporal cost reduction.

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

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

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

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

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

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

Original: [Table]
Revised: [Table]

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

Original: [Table]
Revised: [Table]

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

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

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

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

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

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

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

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

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

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

Original: Our second experiment is a computer vision task demonstrating the generalizability of UnfoldML .
Revised: Our second experiment is a computer vision task demonstrating the generalizability of UnfoldML .

Original: Our goal identiﬁcation of images from a subcategory in a dataset with hierarchical label structure.
Revised: Our goal identiﬁcation of images from a subcategory in a dataset with hierarchical label structure.

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

Original: We want to identifyimages with the subcategory from from one of two categories: ‘people’ or ‘vehicles .
Revised: We want to identifyimages with the subcategory from from one of two categories: ‘people’ or ‘vehicles .

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

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

Original: In our experiments Stage-1 predictsthe ‘coarse’ label, and Stage-2 predicts the ‘ﬁne’ label.
Revised: In our experiments Stage-1 predictsthe ‘coarse’ label, and Stage-2 predicts the ‘ﬁne’ label.

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.
Revised: Our model zoo is made of WideResNets (WRN) using Sharpness-Aware Minimization (Foret et al.

Original: ,2020) which is cited as having state-of-the-art performance on the CIFAR-100 image recognition task.
Revised: ,2020) which is cited as having state-of-the-art performance on the CIFAR-100 image recognition task.

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

Original: For
Revised: For

Original: Stage-1 we train multi-class models using full dataset with ‘coarse’ labels as the target.
Revised: Stage-1 we train multi-class models using full dataset with ‘coarse’ labels as the target.

Original: For Stage-2we target the ‘ﬁne’ subcategory that happens-after the ’coarse‘ subcategory from Stage-1.
Revised: For Stage-2we target the ‘ﬁne’ subcategory that happens-after the ’coarse‘ subcategory from Stage-1.

Original: We train
Revised: We train

Original: Stage-2 models using data from the coarse-granularity label.
Revised: Stage-2 models using data from the coarse-granularity label.

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

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

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

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

Original: [Table]
Revised: [Table]

Original: We compare UnfoldML against a baseline multiclass WRN classiﬁer which we train for directlypredicting labels at the ﬁne-granularity.
Revised: We compare UnfoldML against a baseline multiclass WRN classiﬁer which we train for directlypredicting labels at the ﬁne-granularity.

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

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

Original: Our results show that UnfoldML is able to ﬁnd the optimal gating policy for the pipeline and achievea spatio-temporal cost-savings of 6.9X with 1% reduction in accuracy for ‘people’ subcategoryclassiﬁcation; we achieve a cost-savings of 4.7X with 0.4% reduction in accuracy for ‘vehicle’subcategory-of-interest identiﬁcation.
Revised: Our results show that UnfoldML is able to ﬁnd the optimal gating policy for the pipeline and achievea spatio-temporal cost-savings of 6.9X with 1% reduction in accuracy for ‘people’ subcategoryclassiﬁcation; we achieve a cost-savings of 4.7X with 0.4% reduction in accuracy for ‘vehicle’subcategory-of-interest identiﬁcation.

Original: An alternative method is Classiﬁcation with Costly Features (CwCF) (Janisch et al., 2019), whichnavigates trade-off space using Deep Q-Learning while varying cost parameters which restrictsusing features when training models.
Revised: An alternative method is Classiﬁcation with Costly Features (CwCF) (Janisch et al., 2019), whichnavigates trade-off space using Deep Q-Learning while varying cost parameters which restrictsusing features when training models.

Original: CWCF requires a pre-trained cost-unaware classiﬁer, whichis used when no features are restricted.
Revised: CWCF requires a pre-trained cost-unaware classiﬁer, whichis used when no features are restricted.

Original: As a fair baseline against UnfoldML , we used the highestaccuracy WRN in CwCF.
Revised: As a fair baseline against UnfoldML , we used the highestaccuracy WRN in CwCF.

Original: We vary the cost parameter and search for the optimal trade-off point in theaccuracy-cost space.
Revised: We vary the cost parameter and search for the optimal trade-off point in theaccuracy-cost space.

Original: CWCF returns two extreme data points with either poor or high accuracy as the
Revised: CWCF returns two extreme data points with either poor or high accuracy as the

Original: WRN-A shows in Table 4.
Revised: WRN-A shows in Table 4.

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

Original:  Conclusions
Revised:  Conclusions

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

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

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

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

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

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

Original: IDK .
Revised: IDK .

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

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

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

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

Original: UnfoldML and the application of the“happens before” insight generalizes to computer vision tasks with 5x cost savings gained for a mere0.4% drop in accuracy.
Revised: UnfoldML and the application of the“happens before” insight generalizes to computer vision tasks with 5x cost savings gained for a mere0.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.