FlexiFlow boosts machine learning accuracy by switching models dynamically

FlexiFlow: Bandit-based Model Switching in ML Workflows

Machine LearningArtificial Intelligence

Summary

Many machine learning systems use just one model to make predictions, but this can miss chances to improve accuracy because different models work better on different kinds of data. The authors show that switching between models can improve results by up to 23%, but current tools don’t do this automatically. They created FlexiFlow, a system that learns which model to use on the fly by balancing speed and accuracy with a smart decision method called a multi-armed bandit. Their system also saves time by reusing partial results and works well for complex real-world tasks.

What this means in practice

  • For ml engineering teams: Improve inference accuracy and efficiency by automatically choosing the best model at runtime within complex ML pipelines.
  • For cloud service operators: Enhance cloud-hosted ML services by integrating dynamic model switching to optimize both prediction quality and compute resource use.$Commercial implications: Enables selling more efficient and accurate ML inference services to enterprise customers with variable workloads.

Authors

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

Abstract

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.