Selecting a Model for ML Systems
In a time-constrained and pressurized environment like an ML system design interview, selecting the right model can feel intimidating. In this lesson, we’ll break down how to pick an appropriate model so that you can:
- Design a working model architecture that addresses the problem statement.
- Build a model pipeline that serves as a baseline for further improvements.
- Anticipate realistic problems faced in an ML pipeline.
Overview of common models
Consider the potential models you can choose, as well as their pros and cons. For a quick refresher, check out “A Tour of Machine Learning Algorithms” and “6 Natural Language Processing Models you should know.”
Once you’ve established a foundational understanding of different models, consider these 4 factors when selecting a model for your system:
- Type of problem
- Use case
- Parsimony
- Practical constraints
Below, we specify how you can incorporate these factors to pick the best model for your system.
Type of problem
For this part of the interview, it’s useful to have a broad familiarity with the standard kinds of models used for different types of learning problems. If your problem is a recommendation problem, for example, consider picking a model for collaborative or user-based filtering. These techniques are well-studied, commonly used for this kind of learning problem, and have passed the test of time.
Use case
Is this a model that will make predictions ingested by another system, or will users directly interact with the model? If the latter, consider a generative model that produces “fluent” text or images. Additionally, more robust evaluations of the safety of the model outputs may be needed. If the former, choose a simpler model that produces numeric outputs, which are post-processed and used by another system (e.g. using predictions of how likely a user will watch a particular movie to rank movies on a streaming video website).
If the model needs frequent re-training, adaptation, or personalization, it’s helpful to select a model that can be made modular or trained very efficiently. Many techniques exist for either efficient training (e.g. LoRA, lower-precision training) or adapting a model for different specialized purposes (e.g. adapters, prompt-tuning, prefix-tuning, fine-tuning).
Parsimony
It’s generally a good idea to pick the simplest possible model that achieves sufficient accuracy. Simpler models are easier to work with, understand, and train. A "simple" model may not just refer to a model with fewer features or lower complexity. For example, linear models and decision trees are often considered to be “simple” models, due to their ease of interpretation and conceptual simplicity to build. On the other hand, deep learning models are considered to be complex, because they can take a long run time to train and they are often difficult to interpret.
Practical constraints
The product you’re trying to deploy this model for may also have a variety of practical constraints. For example:
- In high-stakes settings (e.g. healthcare, self-driving cars), it may be desirable to use a more interpretable and higher precision model.
- In compute-constrained settings (e.g. on-device applications), it may be important to select an efficient/lower-capacity model or one that is easily compressed.
- If the product will occasionally experience traffic spikes, it may be helpful to temporarily switch to a slightly lower accuracy (but more efficient) model to ensure the availability of model predictions.
- For time-sensitive cases, an off-the-shelf model that can be deployed in a matter of hours may be the optimal decision. In contrast, most industry problems need to be re-trained on custom datasets, which take a longer timeframe.
- Lower-resourced companies like startups may not have the funds or personnel to train large models or crowdsource their labeling tasks for their datasets. Choosing pre-trained models or using self-supervised learning to deal with unlabeled data may work better for companies with budget constraints.
Now that we’ve reviewed best practices for model selection, we’ll cover model evaluation techniques in Evaluating an ML Model.