[Notes] Designing ML Systems

Overview of ML Systems

Latency isn’t an individual number but a distribution. Think in percentiles (e.g. 50th percentile, 99th percentile, etc..). Highest percentiles can be the most impt users

Intro to ML Systems Design

Crucial to tie the performance of an ML system to the overall business performance. What business performance metrics is it influencing? E.g. amt of ad revenue

ML systems should be

• Reliable: Perform at the correct function at the desired level of performance
• Scalability: Resource scaling, artifact management, automate the monitoring & retraining aspect
• Maintainability: Structure workload so different contributors can work with their own tools. Code, data, artifacts versioned; models sufficiently reproducible
• Adaptability: To adapt to shifting data distributions & business requirements, the system should have capacity to discover aspects for performance improvement & allow updates without service interruption

Minimizing multiple objective functions

1. Combine both losses into one loss & train a model to minimize this loss. Pareto optimization deals with optimization involving >1 objective function
2. Train two different models. Use their combined scores

If there’s multiple objectives, best to decouple them first - easier to maintain (diff objectives have diff maintenance schedules), and easier to tweak each individual model

Data Engineering Fundamentals

Processing sequential data is more efficient than non-sequential data: if a table is row-major, accessing its rows will be faster than accessing its columns

• Row-major: Consecutive elements in a row stored next to each other in memory. Better for accessing examples (e.g. CSV)
• Column-major: Consecutive elements in a col stored next to each other in memory. Better for accessing features (e.g. parquet)

Pandas is built around DataFrame, which is column-major. A NumPy ndarray is row-major by default if we don’t specify the order.

Main ways to represent data

• Relational: Require our data to follow a strict format
• Non-relational: Main ones include document model (self-contained docs, relationships are rare), graph model (relationship between data items are common & impt)

Two types of workloads databases are optimized for

1. Transactional processing: Inserted as they are generated. Need to be processed fast (low latency) so they don’t keep users waiting.
2. Analytical processing: Aggregating data in columns across multiple rows of data. Efficient with queries that require you to look at data from different viewpoints

3 main modes of dataflow

1. Passing through databases: To pass data from A to B, A can write to the db, and B simply reads from that db. Both processes must be able to access the same database
2. Passing through services: Send data directly through a network that connects both processes. Most popular styles are REST & RPC

• Request-driven: To pass data from B to A. A sends a request to B that specifies the data A needs. B returns the requested data. This is synchronous
• Service: A process that can be accessed remotely. E.g. expose B as a service to A
• Microservices: Two services communicating with each other are part of the same application

3. Passing through real-time transport: Use in-memory storage to broker data. Most common types are pubsub, message queues

• Event-driven: A piece of data broadcast to a real-time transport is called an event
• Pubsub: Any service can publish to different topics, and subscribe to a topic to read all events in that topic. E.g. Apache Kafka, Amazon Kinesis
• Message queue: Event often has intended consumers. E.g. RocketMQ, RabbitMQ

Batch processing: When data is processed in batch jobs. Distributed systems like MapReduce & Spark are batch computation engines that process batch data efficiently

• Used to compute features that change less often, aka static features

Stream Processing: Computation on streaming data (data in real-time transports like Kafka). Streaming computation engines include Apache Flink, Spark Streaming

• Used to compute features that change quickly, aka dynamic features

Training Data

Sampling is necessary for creating training data - we don’t have access to all possible data so we often train with a subset of real-world data, and it may be infeasible to process all the data we have access to as it’s time/resource intensive.

Two types of sampling

• Nonprobability sampling: Not based on any probability criteria - it’s convenient but riddled with selection biases. Convenience sampling, snowball sampling (selected based on existing samples), judgment sampling (experts decide), quota sampling (select samples based on quotas for certain slices of data)
• Probability sampling: Simple random sampling (equal probability of being selected; rare categories might not appear), stratified sampling (divide population into groups, sample from each group separately), weighted sampling (each sample given a weight that determines probability of it being selected), reservoir sampling (ensure streamed data has equal prob of being selected), importance sampling (approximate expectations under difficult-to-sample target using a simpler-to-sample distribution)

Most ML models are supervised, meaning they need labeled data. Two methods of labeling - hand-labeling, natural labeling

• Hand-labeling: Can be expensive & time-consuming. If task/data changes (e.g. adding a new class), data needs to be relabeled. Multiple annotators are used, leading to different levels of accuracy.
• Natural labeling: The model’s predictions can be automatically evaluated (e.g. predict stock prices.. Can see the actual price in the near future)

Can set up a system in a way that allows you to collect some feedback/natural labels. E.g. Newsfeed, can add buttons & other reactions to each item to collect feedback. Google translate, can allow users to submit alternative translations for bad translations

Techniques have been developed to address challenge of acquiring sufficient high-quality labels

• Weak supervision: Label data with heuristics (rules-based). Uses a labeling function that encodes this heuristics. Can be noisy. E.g. Snorkel
• Semi-supervision: Generate new labels based on a small set of initial labels. E.g. self-training (train model on existing data, and use it to make predictions for unlabeled samples. Train new model using this expanded dataset), perturbation-based (apply small adjustments to data, e.g. white noise to img)
• Transfer Learning: A base model is reused as a starting point for another task. Using pretrained model as a starting point often boost performance significantly
• Active Learning: Models choose most helpful unlabeled samples for annotators to label. E.g. uncertainty measurement (label samples the model is least certain about lowest probabilities for the predicted class), query-by-committee (label samples that an ensemble of models disagree on the most)

Class imbalance is the norm. 3 ways to handle: choose right metrics, data-level (change data distribution to make it less imbalance), algo-level (use more robust learning method)

• Eval metrics: ROC curve (plots true positive against false positive rate for different confidence thresholds)
• Data-level: Resampling (oversample minority, undersample majority). SMOTE & Tomek Links are effective undersampling methods for low-dimensional data. Two-phase learning where model first trains on a balanced dataset (undersample majority), and then fine-tune on original data (to recover info from majority)
• Algo-level: Most involve adjusting the loss function, which guides the learning process. Cost-sensitive learning (misclassifying different classes incur different costs. Have a cost matrix to specify the cost if class A is misclassified as B). Class-balanced loss (make the weight of each class inversely proportional to no. of samples - rarer classes have higher weights). Focal loss (focus on learning samples it has difficulty classifying. Give more weight if the sample had lower probability of being right)

Data augmentation is used to increase amt of training data. 3 main types: simple label-preserving transformations, perturbation, data synthesis

• Label-perserving transformations: For CV, can crop, flip, rotate, invert image. For NLP, randomly replace a word with a similar one
• Perturbation: Add a small amount of noise to data, which helps the model improve decision boundaries. Can add via injecting random noise, or by search strategy. DeepFool finds minimum possible noise injection to cause misclassification. BERT replaces a small amt of chosen tokens with random words, shown to give small boost in performance
• Data Synthesis: Templates are a cheap way to bootstrap model if you have list of data (e.g. Find me a {CUISINE} within {DIST} miles). Mixup trains models on linear combinations of inputs & their labels (e.g. mix an image of a dog & cat into one by taking their weighted avg at pixel level, and get a soft label like 30% dog, 70% cat)

Model Development, Offline Eval

Simple way to estimate how model performance might change with more data is to use learning curves (plot performance against no. of training samples used)

Ensembles consistently give a performance boost. Each model in the ensemble is a base learner; the less correlation there is among base learners, the better the ensemble

• Bagging: Train multiple models, each on a bootstrap (subset) of the dataset. Use the combined outputs of these models (e.g. majority/average)
• Boosting: Multiple models are trained on the same dataset. For each model, previously misclassified samples are given higher weight (gives a higher loss). Classifiers with lower training error are given more weight to the final prediction
• Stacking: Create a meta-learner that combines the outputs of the base learners to output final predictions

ML systems are part code, part data - so we need to version both!

Two main types of parallelism

• Data parallelism: Split data on multiple machines, accumulate gradients. 2 approaches - synchronous SGD (prone to stragglers slowing down entire system), asynchronous SGD (might lead to gradient staleness where worker is using outdated weights)
• Model parallelism: Different components of the model is trained on different machines (split by tensor/layer, etc.)

Evaluate your model against baselines

• Random Baseline: Expected performance if our model just predicts at random
• Simple heuristic: Make predictions based on simple heuristics (e.g. classify spam based on backlist)
• Zero rule baseline: When your baseline model always predicts the most common class (e.g. recommend the most frequently used app)
• Human baseline: How model performs compared to human experts
• Existing solutions: How model compares to existing solutions (e.g. business logic)

Evaluation methods - we want our model to be robust, fair, and calibrated

• Perturbation tests: Inputs in production are often noisy. Make small changes (e.g. add background noise) to your test splits to see how they perform with noisy data
• Directional expectation: Certain changes to inputs should cause predictable changes in output (E.g. increase in housing size -> increase in price)
• Model calibration: A model is calibrated if its predicted probabilities match real-world frequencies. A simple method is plot the number of times the model outputs the probability X and the frequency Y of that prediction coming; perfectly calibrated model will have x = y
• Slice-based evaluation: Simpson’s paradox when trends appear in individual groups but disappear/reverse when groups are combined - aggregation conceals/contradicts actual situations. To evaluate performance, must eval not only on entire data, but on individual slices. Slices can be found via heuristics-based (using domain knowledge), error analysis (go through misclassified examples & find patterns), slice finder (find slices using algos)

Model Deployment & Prediction Service

Exposing an endpoint for others to use your model is the easy part. The hard part is making your model available to millions of users with a latency of milliseconds and 99% uptime, setting up proper error notification channels, and seamlessly deploying updates.

ML deployment myths

• You only deploy 1-2 ML models at a time: An app may have many different features, with each feature requiring its own model. Moreover, each country may need its own set of models until you have models that can generalize across different cultures & languages
• Model performance remains the same: Due to data distribution shifts, the model performs best right after training and degrade over time
• You won’t need to update your models as much: As our model’s performance decay over time, we want to update it as fast as possible

3 main modes of prediction

• Batch prediction: Generated periodically or when triggered. Uses features computed from historical data. E.g. item embeddings
• Online prediction using batch features
• Online prediction using batch + streaming features: Streaming features are computed from data in real-time transport

Batch prediction is largely legacy - most companies have a batch system to make predictions. When they want to use streaming features for online prediction, they need to build a separate streaming pipeline.

• Building infrastructure to unify stream processing and batch processing has become a popular topic in recent years for the ML community

3 ways to reduce inference latency - make the algo do inference faster, make the model smaller (model compression), make the hardware it’s deployed on run faster.

4 main techniques for model compression

• Low-rank optimization: Replace high dimensional tensors with a product of lower-dimensional ones
• Knowledge-distillation: Small model trained to mimic a larger model. Can work regardless of architectural differences between teacher & student (e.g. random forest as student, transformer as teacher)
• Pruning: Remove entire nodes of a neural network (reduce no. of params), or set least useful params to 0
• Quantization: Use fewer bits to represent parameters. Can happen during training (quantization aware training) or post-training (quantized for inference)

Edge computing is ideal as cloud services can be costly, and edge allows you to run apps in places without internet connections. Less network latency as data doesn’t have to be transferred to the cloud. Can handle sensitive user data. The edge devices must be powerful enough to handle the computation, have enough memory to store ML models, and enough battery (or connected to an energy source) to power the app.

For models built with a certain framework (e.g. PyTorch), that framework must be supported by the hardware vendor. Different hardware backends have different memory layouts (L1, L2, L3 layouts) & compute primitives (1D vector for GPU, 2D matrix for PTU, etc.) - performing an operation on each hardware will be different.

• Intermediate representation (IR) are structured representations of a program that compilers use to bridge frameworks with hardware platforms. Compilers generate high-level (computation graphs, etc.) & low-level IRs (LLVM, PTX, etc.) so the code can run on that hardware backend

4 common ML model local optimizations (optimize a set of operators)

• Vectorization: Execute multiple elements contiguous in memory rather than one-by-one (for loops) to reduce latency caused by data I/O
• Parallelization: Divide input array into different, independent work chunks, and do the operation on each chunk individually
• Loop tiling: Change the data accessing order in a loop to leverage hardware’s memory layout and cache. Optimization is hardware dependent. A good access pattern on CPUs is not a good access pattern on GPUs.
• Operator fusion: Fuse multiple operators into one to avoid redundant memory access

AutoTVM tunes individual operators (kernels) by searching over implementation schedules (how a computation is executed) to find the fastest one for a given hardware target.

To run executable programs (e.g. ML code) in browsers, we can compile our model to WebAssembly (WASM) format, getting an executable file that we can use with javascript. Although WASM is faster than JavaScript, it’s much slower compared to running natively on devices (iOs/android apps)

Data Distribution Shifts & Monitoring

A model’s performance degrades over time. There are 2 main types of failure: software system failures & ML-specific failures

Software system failures

• Dependency failure: A dependency (codebase, package) your system depends on breaks, leading your system to break
• Deployment failure: Caused by deployment errors (e.g. lack read/write permissions for certain files)
• Hardware failure: Hardware to deploy your model doesn’t perform as intended (e.g. CPU overheats & breaks down)
• Downtime/Crashing: A component of the system is down (e.g. AWS)

ML-specific failures

• Production differs from training data: Real world isn’t stationary & virtually infinite, unlike training data which is static & limited. Unlikely for training data to be completely same as the underlying distribution of the training data
• Edge cases: Samples so extreme they cause the model to make catastrophic mistakes, especially for safety-critical applications (e.g. autonomous vehicles)
• Degenerate feedback loops: System outputs are used to generate system’s future inputs, influencing future outputs - causing a system’s outputs to be more homogenous over time. Detect such feedback loops by measuring popularity diversity of a system’s outputs - low score = outputs are homogenous, not good. Can correct such loops by introducing randomization, or encode positional features

3 subtypes of data distribution shifts

• Covariate shift: Input distribution shifts while prediction function stays fixed. If you know in advance how real-world input distribution differs from your training data, use importance weighting - estimate density ratio between real-world input distrib & training input distrib, and weight the training data based on this ratio
• Label shift: Output distribution changes, but the input distribution P(X|Y) for a given output is the same
• Concept drift: Rules for making predictions have changed i.e. P(Y|X) changes but P(X) remains the same

We can monitor distributions to see if they’ve changed

• Simple methods compare their statistics like quartiles, min, max, median, etc.
• Two-sample hypothesis test determines whether the difference between two populations (datasets) is statistically significant. E.g. Kolmogorov-Smirnov test (KS test), Least-Squares Density Difference, Maximum Mean Discrepancy

Address data distribution shifts via

1. Train models with massive datasets so they can learn a comprehensive distribution
2. Adapt a trained model to a target distribution without requiring new labels. E.g. use causal interpretations together with kernel embedding of conditional and marginal distributions to correct models’ predictions
3. Retrain your model using labeled data from the target distribution

If you want to retrain your model, 2 qns: whether to train from scratch (stateless retraining) or continue training from last checkpoint (stateful training), and what data to use - last 24 hours, last week, 6 months, etc..

Can design your system to make it more robust to shifts. Consider the trade-off between the performance & stability of a feature: might be really good for accuracy but deteriorate quickly, forcing you to train your model more often

Monitoring (tracking so we know when something goes wrong) & observability (what went wrong)

Monitoring is about metrics

• Operational metrics (network, machine, application - e.g. latency, throughput, % of successful requests, CPU/GPU utilization)
• ML-specific metrics (accuracy-related, predictions, features, raw inputs) - monitor predictions for distribution shifts, and monitor features to ensure they follow the expected schema (ranges, median, etc.)

Monitoring toolbox include: logs (best to use stream processing rather than batch process so you can discover anomalies in your logs asap), dashboard (visualize metrics so it’s easier to monitor), alerts (notify the right people when things go wrong)

Continual Learning & Test in Production

Companies that employ continual learning in production update their models in micro-batches (e.g. after every 512/1024 examples)

Most companies do stateless retraining - the model is trained from scratch each time. Stateful training (aka fine-tuning or incremental learning) means the model continues training on new data.

There are 2 types of model updates

• Model iteration: New feature is added to an existing model architecture, or the model architecture is changed
• Data iteration: Model architecture & features remain the same, but you refresh this model with new data

Continual Learning Challenges

• Fresh Data access: The speed at which a model can be updated is bottlenecked by the speed at which data is labeled. Can leverage programmatic labelling tools like Snorkel to generate fast labels
• Evaluation: More frequent model updates = more opportunities for failure
• Algorithm: Neural networks are easier to adapt, compared to matrix-based (e.g. collaborative filtering, have to build the entire user-item matrix again)

4 stages of continual learning

1. Manual, stateless retraining: Process of updating a model is manual and ad hoc
2. Automated retraining: Write a script to automatically execute all retraining steps. Some run experiments to determine the optimal retraining frequency.

• When creating automation scripts, must account for different models in your system that might require different retraining schedules
• More complicated if there are dependencies among models (e.g. if model A depends on B, then if B is updated, A should be updated too)
• 3 major factors that affect feasibility of automation script - (a) scheduler (handles task scheduling, e.g. Airflow, Argo), (b) availability of data (need from multiple orgs? Need to label?), (c) model store to automatically version & store all artifacts needed to reproduce a model

3. Automated, stateful training: Update automatic updating script so when we update the model, it first locates the previous checkpoint and loads it into memory before continuing training on this checkpoint/ Need a way to track our data, model lineage - so we know how our models evolve overtime
4. Continual Learning: At stage 3, models are still updated based on a fixed schedule. Instead, you might want to automatically update models whenever data distributions shift and the model’s performance plummets
   1. Trigger mechanism can be time based (every 5 min), performance-based (when model performance plummets), volume-based (when amt of labelled data increases by 5%), drift-based (when major data distribution is detected, etc.)

Offline evaluation isn’t enough. Test splits are usually static, which may not be representative of current data. The only way to know if a model will do well in prod is to deploy it - test in production

• Shadow deployment: Deploy candidate model in parallel with existing model. For each incoming request, route it to both models, but only serve the existing model’s prediction to the user. Log predictions from the new model for analysis purposes.
  • Only replace the existing model when you find the new model’s predictions are satisfactory. Is expensive as you double the no. of predictions generated
• A/B testing: Deploy the candidate model alongside the existing model. A percentage of traffic is routed to the new model for predictions; the rest is routed to the existing model for predictions. Monitor and analyze the predictions and user feedback to see the models performances
  • There are cases where one model’s predictions might affect another model’s predictions e.g., in ridesharing’s dynamic pricing, a model’s predicted prices might influence the number of available drivers and riders, which influence the other model’s predictions. You might have to run your variants alternatively - serve model A one day, serve model B the next day
  • Traffic routed to each model must be truly random
  • A/B test should be run on a sufficient number of samples to gain enough confidence about the outcome
• Canary Release: Slowly roll out the change to a small subset of users, before making it available to everybody. A portion of traffic is routed to the candidate model, while the rest is routed to the existing model. If performance is satisfactory, increase traffic to the candidate model, until the candidate model has replaced the existing model
  • Don’t need to randomize traffic routing. Can roll out the candidate model to a less critical market first
• Interleaving Experiments: Two algorithms can be compared by measuring user preferences (e.g. expose users to recommendations from both models and see which recommendation they clicked on)
  • Team-draft interleaving: The position of a recommendation influences how likely a user clicks on it. We need to ensure a recommendation is equally likely to be generated by A or B - for each recommendation, randomly select A or B with equal probability, and pick the top recommendation that hasn’t already been picked
• Bandit: Determine how to route traffic to each model while maximizing prediction accuracy. Stateful before routing a request to a model, you need to calculate all models’ current performance. You need to get feedback on whether a prediction is good or not (e.g. feedback from user) so you can update the payoff of each model quickly.

Infrastructure & Tooling for MLOps

4 main layers of infrastructure

• Storage & compute: Storage layer is where data is collected & stored. Compute layer provides the compute needed to run ML workloads such as training a model, computing features
  • Compute layer can be sliced into smaller compute units (e.g. CPU can be sliced into threads). When evaluating a new compute unit, evaluate how long it will take this CU to do common workloads. MLPerf is a popular benchmark measuring how fast it takes hardware to train ResNet-50 model
  • A production service might be spread out on multiple instances. No. of instances changes from time to time. Container tech (Docker) helps us re-create an environment on any new instance, dynamically allocating instances as needed. K8s (Kubernetes) helps us create a network for container to communicate & share resources
• Resource management: Tools to schedule & orchestrate your workloads to make the most out of your available compute resources. E.g. Airlow, Kubeflow
  • Two key characteristics that influence their resource management: repetitiveness (e.g. cronjob) and dependencies (deal with workflows, e.g. if A fails then run B)
  • Slurm is a popular scheduler you specify the job name, time when the job needs to be executed, and amount of memory and CPUs to be allocated for the job
  • Orchestrators are concerned with where to get those resources. Deal with lower-level abstractions like machines, instances, etc…
  • Workflow management tools manage workflows - typically specify workflows as DAGs. Most common ones are Airflow, Argo, Prefect, Kubeflow (define workflow in Python, need a Dockerfile & YAML file to specify specs of each component), Metaflow (fully python)
• ML Platform: Aid the development of ML apps like model stores, feature stores, and monitoring tools
  • Deployment service help push your models to production & expose them as endpoints
  • Model store doesn’t just store model itself, but the info associated with it (model parameters, featurize/predict functions, dependencies, data, model generation code, etc.)
  • Feature store solves 3 issues: feature management (share features which multiple models can use), feature computation (save the results after a feature is computed, to save compute), and feature consistency (ensure consistency between features in training & inference)
• Dev Environment: Where code is written and experiments are run. Code needs to be versioned and tested; experiments need to be tracked