Source-linked AI summary
KeystoneML: Optimizing Pipelines for Large-Scale Advanced Analytics
Evan R. Sparks, Shivaram Venkataraman, Tomer Kaftan, Michael J. Franklin, Benjamin Recht
TL;DR
Large-scale ML applications require complex, resource-intensive pipelines, while existing systems provide limited support for constructing and optimizing them end to end. KeystoneML captures these applications through a high-level API and optimizes execution across operators and whole pipelines in distributed environments. Across real workloads, it achieved up to 15× higher performance than unoptimized execution while maintaining scalable and strong statistical performance.
Problem
Existing systems provide little support for automatically constructing and optimizing complex, multi-stage ML pipelines spanning feature processing and model training.
Method
KeystoneML uses high-level logical operators to capture end-to-end ML applications and applies operator-level and whole-pipeline optimizations in distributed execution.
Results
Up to 15× performance improvement over unoptimized execution was achieved through combined physical-operator and end-to-end optimization across real-world ML workloads.
Takeaways & Limitations
KeystoneML provides scalable, high-throughput training with statistical performance comparable to recent results across image, speech, and language applications.
Takeaways & Limitations
Runtime estimates are imperfect, so KeystoneML falls back to LRU cache replacement when estimates are inaccurate.
Abstract
from arXiv · showhide
Modern advanced analytics applications make use of machine learning techniques and contain multiple steps of domain-specific and general-purpose processing with high resource requirements. We present KeystoneML, a system that captures and optimizes the end-to-end large-scale machine learning applications for high-throughput training in a distributed environment with a high-level API. This approach offers increased ease of use and higher performance over existing systems for large scale learning. We demonstrate the effectiveness of KeystoneML in achieving high quality statistical accuracy and scalable training using real world datasets in several domains. By optimizing execution KeystoneML achieves up to 15x training throughput over unoptimized execution on a real image classification application.
1 Introduction
KeystoneML addresses the lack of end-to-end support for constructing and optimizing complex ML pipelines by providing a high-level system that optimizes operators and whole workflows. Its evaluations show substantial speedups, scalability, and strong statistical performance across real-world domains.
- Advanced analytics applications combine feature extraction, transformations, dimensionality reduction, and supervised learning, but existing systems provide little automated pipeline construction or optimization.
- Developers commonly assemble domain-specific libraries and general-purpose learning packages, creating cumbersome, error-prone pipelines that require re-engineering as data or features grow.
- KeystoneML lets users specify end-to-end ML applications with high-level logical operators, scale them as data and complexity change, and optimize them automatically.
- Its cost-based optimizer accounts for computation and communication costs, while a greedy algorithm selects intermediate states to materialize during iterative execution.
- 7× end-to-end optimization speedup and up to 15× speedup with physical-operator and end-to-end optimization were observed versus unoptimized execution.
- Evaluations across computer vision, speech, and natural language workloads showed near-linear scalability over hundreds of machines and strong statistical performance.
2 Pipeline Construction and Core APIs
KeystoneML represents end-to-end machine-learning applications as high-level pipelines of logical operators, then optimizes their execution as a distributed operator DAG. Its API supports composable Transformers and Estimators, branching, and operator-level optimization.
- Logical operators: KeystoneML uses high-level logical operators to represent complete machine-learning applications instead of piecing together imperative libraries.This abstraction is intended to simplify application construction and enable broader optimization.
- Logical operators: Operators with multiple physical implementations can be optimized using associated cost models, while iterative operators are explicitly marked.Logical operators define the computation; physical operators implement it.
- Transformers and Estimators: Transformers apply deterministic, side-effect-free functions to individual items or collections, enabling execution reordering without changing results.Examples include data transformations, feature extractors, and model application.
- Transformers and Estimators: Estimators consume distributed data and produce Transformers; for example, LinearSolver learns a model from data and labels for later application.This separates model fitting from applying the learned model to new data.
- Pipeline composition: Pipelines support linear chaining and branching through andThen and gather, while common-subexpression optimization removes redundant computation.The API incrementally builds composable pipeline structures.
- Pipeline execution: At optimization time, KeystoneML builds and optimizes a pipeline DAG, then executes it depth-first with operators up to Estimator pipeline breakers packed into jobs.Lazy optimization gives the optimizer full information about the application before runtime execution.
3 Operator-Level Optimization
KeystoneML selects physical implementations for distributed ML operators using cost models that incorporate computation, communication, input statistics, and cluster resources. Experiments show that operator choices vary across workloads and can substantially affect runtime.
- Cost-based operator selection: KeystoneML’s operator-level optimizer chooses physical implementations using computation and communication costs for distributed ML operators.The cost model is designed to accommodate new operators and hardware configurations.
- Cost-based operator selection: The cost model uses dataset statistics such as sparsity and dimensionality alongside cluster resources including CPU, memory, storage, and network capabilities.Resource information is collected through configuration data and microbenchmarks.
- Linear solvers: 5-20× faster than the exact solver, L-BFGS outperformed alternative solvers on sparse Amazon text features as feature counts increased from 1k to 16k.The exact solver exceeded memory capacity above 4k features, while L-BFGS exploited sparsity for cheap gradient calculations.
- Linear solvers: 3-9× faster than L-BFGS, the exact solver was better for smaller dense TIMIT problems, but the block-wise solver became preferable beyond 8k features.The results show that solver choice changes with both problem size and input sparsity.
- Evaluation across operators: 84% of PCA choices matched empirical best selections, while the optimizer selected the right linear solver 90% of the time.For PCA, increasing data volume favored distributed execution, while small k could favor approximate PCA; convolution choices also changed with filter size.
- Evaluation across operators: Convolution implementation depends on filter size and separability: BLAS is fastest for small filters, matrix-vector methods suit separable filters, and FFT performance is independent of filter size.The alternatives include matrix-matrix multiplication, separable matrix-vector products, and FFT-based convolution.
4 Whole-Pipeline Optimization
KeystoneML optimizes whole ML pipelines by using data statistics, computation reuse, and memory-aware materialization. Its caching formulation models execution time under memory constraints and motivates a practical greedy algorithm.
- Operator statistics: KeystoneML collects input statistics, such as vector sparsity, so downstream operators can select suitable optimizations.Statistics are gathered at each pipeline stage because later decisions depend on transformed-data properties.
- Common sub-expression elimination: Common sub-expression elimination merges repeated pipeline computations, allowing shared training data or featurization outputs to be reused.The text-classification example reuses document bigrams across feature selection and classifier training.
- Automatic materialization: Materializing intermediate results can reduce repeated work because iterative pipelines may revisit the same data multiple times.KeystoneML captures pipeline iteration in a DAG to identify reuse opportunities.
- Automatic materialization: Intermediate outputs may reach multiple terabytes, making unconstrained caching and policies such as LRU potentially suboptimal.Caching a large intermediate feature object can evict a smaller, expensive-to-recompute object.
- Automatic materialization: KeystoneML selects cache items using access frequency, output size, and materialization runtime under a memory budget.The objective is to choose cached nodes that minimize total execution time.
- Automatic materialization: The formal cache schedule uses binary cache indicators and DAG relationships to estimate execution time, while a greedy algorithm avoids impractical ILP-solving costs.The greedy method chooses nodes with the greatest execution-time savings subject to available memory.
5 Evaluation
KeystoneML evaluates end-to-end machine-learning pipelines across multiple domains and compares scalability, optimization strategies, and competing systems. Its optimizations improve throughput substantially, while scaling is limited by pipeline stages such as featurization and solving.
- Workloads: KeystoneML evaluates text, image, and speech pipelines on real-world workloads, comparing performance and statistical accuracy with prior implementations.The experiments use distributed clusters and aim for comparable or better statistical results than the respective benchmark authors.
- Workload characteristics: Intermediate datasets can grow by orders of magnitude during featurization before entering the solver, potentially exceeding available cluster memory.The evaluation documents this growth in its dataset characteristics and notes TIMIT as an example where the solve input is too large for cluster memory.
- Overall performance: KeystoneML achieves one to two orders of magnitude higher end-to-end throughput than a single node and equivalent or better performance than comparable cluster systems.The authors associate these improvements with quicker ML application development and higher developer productivity.
- System comparison: 5.5× faster training is achieved at 65536 features than SystemML, which converges to worse training loss over 10 iterations.KeystoneML finishes in 17 minutes, while SystemML takes 1 hour and 40 minutes.
- System comparison: KeystoneML surpasses TensorFlow on the compared workload with a minimum runtime of 29 minutes, versus TensorFlow’s 57 minutes.KeystoneML continues improving through 32 nodes, whereas TensorFlow achieves its best result on a 4-node cluster.
- Optimization breakdown: Whole-pipeline optimization improves the Amazon pipeline by 7×, while operator optimization adds no further benefit because the default L-BFGS solver is already optimal.The gains come from caching intermediate features, and operator optimizations produce a 12× improvement on another breakdown, or 15× when optimization costs are amortized.
- Materialization: The materialization strategy is nearly always better than rule-based caching or LRU, while estimated runtimes are within 15% of actual runtimes.The strategy adapts to constrained and unconstrained memory, and inaccurate estimates fall back to LRU replacement.
- Scalability: Scaling is sub-linear when a pipeline stage dominates and requires coordination: solving dominates TIMIT, while featurization dominates Amazon and ImageNet.The authors identify coordination requirements in linear solvers as a direct source of sub-linear whole-pipeline scalability.
6 Related Work
Related systems address portions of large-scale machine learning, but KeystoneML emphasizes end-to-end pipelines that include featurization and operator-level optimization. Its approach extends optimization beyond conventional relational assumptions and supports iterative, approximate, and data-dependent ML computation.
- ML frameworks: Existing ML libraries largely optimize individual learning algorithms, whereas KeystoneML focuses on pipelines that combine featurization with end-to-end performance optimization.The paper distinguishes its scope from systems centered on regression, classification, and linear algebra routines.
- Optimization systems: SystemML uses optimization to select physical execution strategies, but KeystoneML places greater emphasis on user-defined functions and featurization.KeystoneML uses an extensible API to profile arbitrary nodes and optimize both nodes and whole pipelines.
- Database-integrated ML: Database-integrated ML systems provide modularity, but the paper states that they do not present cross-operator optimizations or operator-level tradeoffs like KeystoneML.KeystoneML targets similar modularity and end-to-end optimization while also supporting scalable execution.
- Optimization scope: KeystoneML handles richer data types and can switch between exact and approximate operators, while traditional relational optimizers often omit sparsity from optimization decisions.Its operators also lack some relational algebra semantics, such as commutativity, limiting certain optimizations.
- Caching: Its caching strategy resembles materialized-view selection but targets intra-query reuse in iterative pipelines rather than inter-query optimization with updates.The approach exploits iterative execution and immutable state.
7 Future Work and Conclusion
The paper presents KeystoneML as a first step toward easy-to-use, robust, and efficient end-to-end ML at massive scale. It concludes that capturing the whole application enables optimization that adapts to data, hardware, and environmental changes, while identifying several directions for future work.
- Future work: Future work includes node reordering, hyperparameter tuning, asynchronous SGD, and back-propagation.The existing APIs are synchronous and current pipelines are acyclic, defining important boundaries for the present system.
- Conclusion: KeystoneML captures complete ML applications and automatically optimizes both individual operators and whole pipelines.The system is intended to adapt execution to changes in data, hardware, and other environmental characteristics.