PaperlayerView exact revision on arXiv

Source-linked AI summary

FlexiFlow: Bandit-based Model Switching in ML Workflows

Abhilash Jindal, Todd Nief, Bhanu Prakash Vangala, Shankaradithyaa V, Tvisha Malik, Anshik Sahu, Aaron Schein, Amitabh Chaudhary, Tanu Malik

arXiv:2610.07286v1cs.LGcs.AI

TL;DR

Soft failures and costly alternative-model testing complicate reliable ML workflows, and existing dataflow systems do not integrate assertions with explore–exploit strategies. FlexiFlow combines model choices, assertions, rollback, and bandit-based selection; experiments report improved accuracy and efficiency across several workflow types.

  • Problem

    Soft failures degrade predictions, and exhaustively testing alternatives is costly; existing ML dataflow systems do not integrate model assertions with explore–exploit strategies.

    Source: §1 · pp. 1–2

  • Method

    FlexiFlow uses Bayesian bandit strategies to select model configurations and rolls back to test alternatives when assertions fail.

    Source: §1 · p. 2

  • Results

    FlexiFlow’s automatic model switching improves overall accuracy and efficiency across document-processing, image-classification, and question-answering pipelines.

    Source: §8 · p. 12

  • Takeaways & Limitations

    FlexiFlow reduces programmer burden by letting users specify operator choices instead of maintaining multiple workflow instances.

    Source: §6 · p. 12

    FlexiFlow is most effective when configuration failures are uncorrelated and does not separately learn conditional assertion-pass probabilities.

    Source: §6 · p. 12

Abstract

from arXiv · show

Model optimizations help improve inference performance and accuracy of ML workflows. However, relying on a single model to perform inference across all data batches often fails to maximize accuracy and thus overall performance. In many cases, alternate models could perform better on specific subsets of data where a primary model underperforms. Our experiments with real ML workflows indeed show that switching models improves workflow accuracy by up to 23%. Yet, current systems lack the ability to adaptively switch between models based on performance, forcing users to manually test models in sequence. We present FlexiFlow, a dataflow system that dynamically switches between alternate models when the current model exhibits low accuracy. FlexiFlow learns to rank models using a novel multi-armed bandit approach that accounts for model runtimes, probability of passing user-defined assertions, and the computational structure of the ML workflow. We show that the standard Thompson sampling approach is insufficient for switching models in ML workflows. In contrast, our proposed approaches are effective and scales to complex real-world ML workflows. Experiments show that switching models at runtime while reusing intermediate results provides higher accuracy, but also 48% efficiency gain compared to sequential workflow runs.

1 Introduction

Soft failures and costly testing make static model choices difficult to sustain across real-world inputs. FlexiFlow combines operator choices, assertions, rollback, and bandit strategies to switch models during workflow execution.

1 Introduction

  • Soft failures can degrade predictions without crashing, while trace-log diagnosis and exhaustive testing of alternatives are resource-intensive and time-consuming.A static model choice may underperform on real-world inputs, forcing the workflow to be rerun from the beginning.Source: §1 · p. 1

1.1 Our Contributions

  • FlexiFlow uses Bayesian strategies to select configurations and, after assertion failures, rolls back to test alternatives without restarting the entire workflow.The system supports multiple model choices at each workflow step.Source: §1 · p. 2
  • FlexiFlow introduces workflow primitives for operator choices and assertions, with adaptive rollback to intermediate operators after assertion failures.Incremental computation reuses prior work after rollback.Source: §1.1 · p. 2
  • Its second hybrid strategy runs bandit optimizations at each workflow step and uses high-payoff configurations to update operator-choice posteriors, providing a learning advantage.The strategy tracks cumulative configuration payoffs.Source: §1.1 · p. 3
  • Adding operator choices improves F1 scores by up to 23% compared with workflows without them, with different configurations often having uncorrelated failures.Strategies must learn operator runtimes and assertion-pass probabilities and account for rollback to find good configurations quickly.Source: §1.1 · p. 3

2 The FlexiFlow Execution Model

FlexiFlow extends a dataflow model with operator choices and workflow assertions. When an assertion fails, the runtime can roll back to an operator and reselect a choice.

2 The FlexiFlow Execution Model

  • A workflow is a DAG of map operators that transform data items, with edges representing data dependencies and schema compatibility checked along them.Operators consume an input schema and produce an output schema; the runtime currently excludes logical operators such as joins and aggregates.Source: §2 · p. 3
  • Each map operator can have alternative implementations, and the runtime selects one at execution time.The operator choice is modeled as a non-deterministic transformation of an input item.Source: §2 · p. 3
  • Users can define Boolean assertions on intermediate or final outputs to detect soft failures and guide operator reselection.Assertions can check semantic properties such as a valid date of birth or confidence above a threshold.Source: §2 · p. 3
  • When an assertion fails, the runtime rolls back to an operator, while a learning-based optimizer guides the rollback and next choice.Rollback can target an intermediate node rather than only the workflow root; the optimizer maintains estimates for operator choices and sets of choices.Source: §2 · pp. 3–4, Figure 4 · p. 4

3 The FlexiFlow Learning Problem

FlexiFlow seeks an ordered sequence of workflow configurations that minimizes expected execution cost while finding a configuration that passes the assertion. Its learning problem accounts for the growing configuration space and the costs of failure and rollback.

3 The FlexiFlow Learning Problem

  • With M operators and k_i choices at operator i, the workflow has K = product of k_i configurations.For the ID-card workflow, 5, 2, and 2 choices yield 20 configurations.Source: §3 · p. 4
  • For each input, FlexiFlow tries configurations in sequence until one passes the assertion, potentially testing all K configurations.The sequence length is at most K, and each additional attempt incurs execution overhead.Source: §3 · p. 4
  • Consecutive configurations sharing a longer prefix of operator choices have lower rollback re-execution cost.The paper illustrates this with two sequences that share the same channel and OCR choices but differ in classifier choice.Source: §3 · p. 4
  • The objective is to find a configuration sequence that minimizes expected runtime, accounting for the chance prior configurations fail and the cost of rollback.Each additional configuration incurs overhead from rerunning operators.Source: §3 · p. 4
  • Because failure probabilities are unknown for new inputs and deployed runtimes can differ from training estimates, FlexiFlow uses online Bayesian explore-exploit learning to identify the sequence.Runtime differences may reflect shifts in data distributions.Source: §3 · p. 4

4 FlexiFlow’s Bayesian Online Optimizer

FlexiFlow’s Bayesian optimizers estimate assertion-success probabilities and runtimes to choose model configurations with low expected runtime. For multi-operator workflows, its strategies account for configuration structure and rollback costs, with a hybrid strategy using operator-level evidence to guide configuration selection.

4 FlexiFlow’s Bayesian Online Optimizer

  • FlexiFlow’s optimizers estimate success probabilities and runtimes to minimize expected execution time.They use Bayesian explore-exploit learning within the multi-armed bandit framework.Source: §4 · p. 4

4.1 FlexiFlow Optimizer for Single Operator with Multiple Choices

  • OTS samples each choice’s assertion-pass probability and runtime, ranks expected time-to-pass, then retries alternatives after failures while updating both estimates.It models pass probability with a Beta distribution and runtime with a log-normal distribution; failed choices are removed as rollback proceeds.Source: §4.1 · p. 5

4.2 FlexiFlow Optimizers for Multiple Operators with Multiple Choices

  • CTS learns at the configuration level and selects a failed configuration’s next alternative by its estimated incremental runtime, while configuration spaces can be large.ID processing requires learning over nearly 500 choices.Source: §4.2 · p. 6
  • HTS combines configuration outcomes with an empirical prior from constituent-operator outcomes, subtracting repeated counts and using γ to control the prior’s influence.This lets operator-level estimates inform configuration-level selection.Source: §4.2 · pp. 6–7
  • OTS models operator-level probability and runtime but does not support multiple-operator, multiple-choice workflows; CTS uses configuration-level probabilities, while HTS also informs them with operator estimates.OTS, CTS, and HTS are equivalent for single-operator workflows with multiple choices.Source: §4.2 · p. 7

5 Evaluation

The evaluation assesses FlexiFlow on synthetic and real ML workflows, examining optimizer runtime and accuracy against individual configurations and how runtime modeling, rollback awareness, and workflow structure affect learning. It also examines operator choices and native rollback support in relation to workflow performance and sequentially running configurations.

5.1.1 Performance of optimizers on synthetic workflows.

  • On synthetic SC, OTS learns the optimal choice order while OTS-R does not; on MC, CTS-RB fails to learn optimal ordering, while HTS converges faster than CTS.Runtime modeling helps OTS account for the tradeoff between assertion success and runtime; rollback awareness helps multi-operator strategies avoid re-executing prior choices, and HTS combines configuration-level learning with operator-level estimates.Source: Figure 8 · p. 9, §5.1.1 · pp. 9–10

5.1.2 Performance of optimizers on real-world ML workflows.

  • On real workflows, OTS converges within 50 runs on IC and SCQA, while HTS runtime converges after about 1,000 runs on LCQA and improves ICP runtimes.OTS favors DeiT-Tiny on IC but RoBERTa on SCQA; HTS delays trying Tesseract until the 17th ICP configuration because it passes assertions nine times less often than EasyOCR.Source: Figure 9 · p. 10, §5.1.2 · p. 10

5.2 Accuracy of FlexiFlow

  • 0.48 F1 on LCQA is FlexiFlow’s score, versus 0.39 for the best single configuration, a relative improvement of 23%.Different configurations make different errors, and FlexiFlow uses estimated confidence to switch configurations when an answer is uncertain. The evaluation uses general-purpose QA models rather than models fine-tuned for HotpotQA.Source: Figure 10 · p. 11, §5.2 · p. 11, Figure 11 · p. 11

5.3 Runtime Performance of FlexiFlow

  • FlexiFlow reduces total runtime by 48% versus sequential workflow-configuration execution in Flink and 49% in Spark, reusing intermediate results during rollbacks and leveraging data parallelism by running different configurations for different inputs simultaneously.Source: §5.3 · p. 11

5.2 Accuracy of FlexiFlow

  • ICP configurations reach a maximum F1 of 0.90, while FlexiFlow operator choices reach 0.94, attributed to uncorrelated failures across configurations.In a QA example, a heavier model returned the correct answer with 0.01 confidence, while a lighter model returned the same answer with 0.98 confidence.Source: §5.2 · p. 11
  • Operator choices using FlexiFlow improve performance by up to 23%.The evaluation attributes accuracy gains to configurations failing on different inputs.Source: §5.2 · p. 11

5.3 Runtime Performance of FlexiFlow

  • With operator choices and assertions disabled and only the best ICP configuration run, FlexiFlow finishes 15% faster than Flink and 22% faster than Spark.This comparison used identical inputs and setup and is shown as FlexiFlow (1) in Figure 12.Source: §5.3 · pp. 11–12, Figure 12 · p. 11

6 Discussion

FlexiFlow lets programmers specify operator choices and uses bandit strategies to reduce the burden of managing multiple workflow instances. The discussion frames this as a first step, with current strategies focused on runtime and future work addressing data-dependent choices, correlated failures, and dollar cost.

6 Discussion

  • FlexiFlow reduces programmer burden by letting them specify operator choices instead of maintaining multiple workflow instances.Bandit strategies also reduce the management burden of prioritizing among workflow instances.Source: §6 · p. 12
  • Future strategies could use input characteristics to learn different optimal configuration sequences for data such as cards from different countries.The paper describes this as future work.Source: §6 · p. 12
  • FlexiFlow is most effective when configuration failure modes are uncorrelated, but its strategies do not separately learn conditional assertion-pass probabilities.The authors note that a heavy model's failure may make a light model's failure more likely, and modeling such conditional probabilities remains future work.Source: §6 · p. 12
  • Current strategies focus on lowering runtimes, while reducing dollar costs in cost-sensitive batch workflows remains a possible future direction.The authors give reducing calls to commercial LLMs as an example.Source: §6 · p. 12

7 Related Work

FlexiFlow extends dataflow workflows with operator choices, assertions, and rollback capabilities that prior systems generally lack. It also differs from related debugging, online-learning, and confidence-estimation approaches in how it selects alternatives and handles workflow errors.

7 Related Work

  • Prior dataflow systems lack alternative operators and assertions, while FlexiFlow can roll back to intermediate outputs unlike systems handling input updates.Timely dataflow supports some rollback but not alternative operator choices.Source: §7 · p. 12
  • Python implementation may contribute to performance differences through Python–JVM interaction overheads for Spark and Flink.Source: §7 · p. 12
  • FlexiFlow applies operator choices in dataflow systems, unlike choice-block programs where an SMT solver selects blocks to meet specifications.The paper describes this as the first application of operator choices in a dataflow system.Source: §7 · p. 12
  • Debugging systems such as Amber and Titian offer inspection primitives, but often require manual changes for correction; FlexiFlow rolls back to alternate workflow configurations.Model assertions [28] detect errors for a single model, whereas FlexiFlow checks errors across workflows.Source: §7 · p. 12
  • FlexiFlow builds on Thompson Sampling [40] but learns online without assuming access to high-quality historical data.The paper reports experimental benefits from learning choices, but says formal regret analysis is beyond its scope.Source: §7 · p. 12
  • Confidence-estimation methods including Bayesian neural networks, ensembles, and temperature scaling do not learn an optimal sequence of models to mitigate ensemble runtime cost.Source: §7 · p. 12

8 Conclusion

The conclusion presents FlexiFlow as a way to improve accuracy and efficiency in ML pipelines through automatic model switching. It points toward adapting workflows to evolving data and incorporating intermediate data characteristics in future learning strategies.

8 Conclusion

  • FlexiFlow combines operator choices, assertions, rollbacks, and a multi-armed bandit approach to improve accuracy and efficiency across document-processing, image-classification, and question-answering pipelines.Source: §8 · p. 12
  • The approach opens avenues for AutoML systems that adapt to evolving data distributions and could reduce manual intervention while enhancing reliability of ML-driven automation.Source: §8 · p. 12
  • Future work will explore incorporating intermediate data characteristics, such as an ID card's country, into learning strategies.Source: §8 · p. 12

Related papers

Loading 2610.07286v1…