ML System Design · Module 12 of 47 · 12 concepts · about 14 min

Choosing and Building the Model

How to pick a model for a production system: start simple, why gradient-boosted trees are the default for tables, when deep learning earns its cost, ensembles and multi-task models, pretrained models, calibrated probabilities, explanations, and improving a model through error analysis.

At work: The step everyone wants to talk about first. In practice it is a choice constrained by the data, the latency budget and the need to explain, made after the framing and the features.

Start simple, then climb

What it is

Climb a ladder of models, one rung at a time, and stop when the gain stops paying:

  • Rules or heuristics: the baseline.
  • Linear or logistic regression: fast, explainable, a strong first model.
  • Gradient-boosted trees: usually the biggest single jump on tabular data.
  • Neural networks: when the data is unstructured, very large, or full of interactions trees cannot capture.

Why it matters

Each rung costs more to train, serve, debug and explain.

  • The simple model shows what the data can give and is the fallback.
  • If boosted trees beat logistic regression by 8% and a neural network beats the trees by 0.5% at five times the serving cost, the trees win.
  • In an interview, naming the ladder and where you would stop shows judgement.

In the real world: Kirana Kart's first ETA model was distance divided by an average speed plus a fixed packing time; the boosted-tree model that replaced it cut error by 40%; a neural network later added 3%.

Visual example

Ladder

Every rung up costs more to run and explain.

Climb only while each step pays for itself

  1. RulescheapestPopular items; simple formulas
  2. Linear modelLogistic or linear regression
  3. Boosted treestabular defaultXGBoost, LightGBM, CatBoost
  4. Neural networkcostliestImages, text, sequences, huge data

Starting at the top. A complex model built first leaves no way to tell whether its accuracy comes from the model or the features, and no simple fallback when it breaks.

Gradient-boosted trees: the tabular default

What it is

Gradient-boosted decision trees (GBDT) build many small trees one after another, each correcting the mistakes of those before. Libraries: XGBoost, LightGBM, CatBoost. Strengths:

  • top accuracy on tables of numbers and categories;
  • little preprocessing: no scaling, missing values handled natively;
  • fast to train and to serve on ordinary CPUs;
  • feature importance and explanations readily available.

Why it matters

For most business prediction problems (fraud, churn, delivery time, pricing, credit) they are the strongest practical choice.

  • Study after study and most data-science competitions on tabular data are won by boosted trees or ensembles of them.
  • Serving hundreds of trees takes well under a millisecond per example.

In the real world: Kirana Kart's fraud, ETA and demand models are all LightGBM; only the product-embedding model behind recommendations is a neural network.

Using boosted trees directly on raw IDs, text or images. They need those turned into features first (counts, embeddings); they cannot learn from pixels or words themselves.

When deep learning earns its place

What it is

Neural networks earn their extra cost when:

  • inputs are unstructured: images, text, audio, video;
  • inputs are sequences: a user's last fifty actions, a time series with long patterns;
  • there are huge numbers of IDs (users, items) best represented as learned embeddings;
  • the data is very large (hundreds of millions of examples) with complex interactions;
  • you need one model for several tasks, or to reuse a pretrained model.

Why it matters

Their power comes with costs.

  • Training needs GPUs and more tuning.
  • Serving is slower and may need GPUs too.
  • They are harder to explain and debug.
  • So use them where their unique strengths matter, often alongside trees rather than instead of them.

In the real world: Large recommendation and ad systems (YouTube, Meta, TikTok) use deep models because they learn from billions of interactions and huge ID spaces, where trees struggle.

Replacing a working boosted-tree model with a neural network on the same thirty tabular features. The usual result is similar accuracy at higher cost.

Comparing models on the constraints

What it is

Choose a model by scoring the options against the requirements written down earlier:

  • accuracy on the right offline measure;
  • latency at p99 for one request (or one batch of candidates);
  • cost to train and to serve;
  • interpretability: can a decision be explained?
  • data needed to train well;
  • operational effort: how hard to retrain, monitor and debug.

Why it matters

A model that wins on accuracy alone can lose on the requirement that actually binds.

  • A table like the one here makes the trade-off explicit for reviewers and interviewers.

In the real world: A lender rules out a neural network for credit decisions despite a small accuracy gain, because every refusal must be explained to the customer with specific reasons.

Visual example

Choice

General tendencies; measure on your own data.

Model families against typical constraints
familytabular accuracylatencyexplainabledata neededbest for
Logistic / linearGoodMicrosecondsYes, directlyLittleBaselines; strict explanations
Boosted treesBestSub-millisecondWith tools (SHAP)ModerateMost business tables
Neural networkSimilar on tablesMilliseconds; may need GPUHardA lotImages, text, sequences, huge ID spaces
Nearest neighbours on embeddingsn/aMilliseconds with an indexBy exampleEmbeddingsRetrieval; similar items

Comparing models only by offline accuracy in a notebook. Latency and serving cost, measured on the real serving stack, often decide.

Ensembles and their production cost

What it is

An ensemble combines several models' predictions:

  • averaging models of different kinds;
  • stacking: a small model learns how to combine the others;
  • boosted trees and random forests are themselves ensembles of trees.

Ensembles usually add a little accuracy.

Why it matters

In production, that little accuracy has a price.

  • Each member must be trained, versioned, served and monitored.
  • Latency is the slowest member's, or the sum if run in sequence.
  • Debugging asks 'which member caused this?'.
  • The famous example: the Netflix Prize-winning ensemble of over a hundred models was never deployed in full, because the gain did not justify the engineering.

In the real world: Kirana Kart keeps a single LightGBM fraud model plus one rules layer, after measuring that a three-model ensemble added 0.2 points of PR-AUC for three times the upkeep.

Shipping a competition-style ensemble. Ten models for 0.3% is almost never worth it outside competitions.

Multi-task models

What it is

A multi-task model predicts several related targets at once, sharing most of its layers:

  • a shared part learns from all targets together;
  • a small head per target outputs each prediction: click, add to basket, purchase.

The final score combines the heads with weights the business chooses.

Why it matters

It is common in ranking systems.

  • Rare targets (purchase) learn from plentiful ones (click) through the shared layers.
  • One model to serve instead of three.
  • The weights between heads let the business trade engagement against revenue without retraining.

In the real world: YouTube described a multi-task ranking model predicting engagement (clicks, watch time) and satisfaction (likes, dismissals) together, combined into one score.

Tasks that conflict. If one target rewards what another punishes, sharing layers can make both worse; check each head against a single-task model.

Pretrained models and fine-tuning

What it is

Instead of training from scratch, start from a model already trained on a large general dataset:

  • use its embeddings as features (product photos through an image model, titles through a text model);
  • fine-tune it: continue training on your own labelled data, often only the last layers;
  • distil it: train a small model to copy a large one, for cheaper serving.

Why it matters

It saves data and time.

  • A pretrained image model already knows edges, shapes and objects; you only teach it your categories.
  • A few thousand labelled examples can then be enough where training from scratch would need millions.

In the real world: Kirana Kart classifies product photos into its categories by training a small classifier on embeddings from a pretrained image model; 20,000 labelled photos were enough.

Fine-tuning on data very different from what the model was trained on, with few examples. It can forget what it knew and fit your small set badly; start by training only a small head on top of frozen embeddings.

Practise it: Design an Image Search System

Calibration

What it is

A model is calibrated if its probabilities mean what they say: of all the cases scored 0.2, about 20% really are positive. Check it by:

  • grouping predictions into bins by score;
  • comparing each bin's average prediction with its observed rate;
  • summarising with the expected calibration error (the weighted average gap).

Fix it with Platt scaling (fit a logistic curve) or isotonic regression on held-out data, after training.

Why it matters

Calibration matters whenever the number itself is used.

  • Expected cost and value ('0.3 chance of fraud × ₹5,000 at risk').
  • Ad auctions, which bid by predicted click probability.
  • Thresholds set in business terms.
  • Combining scores from several models.
  • A model can rank perfectly and still be badly calibrated; ranking measures will not show it.

In the real world: Kirana Kart's risk team computes expected fraud loss per order from the model's score; after resampling made the scores too high, isotonic calibration brought the predicted losses back in line with actual losses.

Visual example

Computed

10,000 fraud scores in four bins. The model is too confident in every bin: expected calibration error 0.069.

Average predicted probability vs the rate actually observed
predictionsaverage predictedobserved rategap
5,0005%2%3%
3,00020%12%8%
1,50050%35%15%
50085%70%15%

Expected calibration error

b0: 5000/10000 × |0.05 − 0.02| = 0.015

b1: 3000/10000 × |0.2 − 0.12| = 0.024

b2: 1500/10000 × |0.5 − 0.35| = 0.0225

b3: 500/10000 × |0.85 − 0.7| = 0.0075

sum: sum = 0.069

Assuming the score is a probability because it is between 0 and 1. Boosted trees, neural networks and any model trained on resampled data are often miscalibrated.

Practise it: Diagnose Model Drift, Design a Large-Scale Classification Platform, Design a Real-Time Fraud Detection System, Design an Ad Ranking System

Explaining predictions

What it is

Two kinds of explanation:

  • Global: which features matter most to the model overall. Feature importance; partial dependence plots.
  • Local: why the model gave this prediction. SHAP values split one prediction into each feature's contribution; LIME fits a simple model around one example.

Why it matters

Explanations serve several people.

  • Engineers use them to debug: a strange top feature often reveals a leak.
  • Reviewers and support staff need reasons for decisions ('flagged because of a new device and an unusual amount').
  • Regulators may require them for credit, insurance and hiring.

In the real world: Kirana Kart's fraud review screen shows the top three SHAP contributions for each flagged order, which roughly halves the time reviewers spend per case.

import shap

explainer = shap.TreeExplainer(fraud_model)
values = explainer.shap_values(order_features)       # one contribution per feature

# Top three reasons for this order's score, for the reviewer's screen
top = sorted(zip(feature_names, values[0]), key=lambda x: -abs(x[1]))[:3]

Reading SHAP values as causes. They describe how the model uses features, not how the world works; a feature can matter to the model because it is correlated with the real cause.

Improve through error analysis

What it is

The fastest way to a better model is to study its mistakes:

  • collect a sample of wrong predictions, especially costly ones;
  • group them by cause: missing feature, bad label, rare segment, data bug;
  • count the groups, and fix the biggest one first;
  • re-check on the slices you fixed, and on everything else.

Why it matters

Error analysis tells you what to do next, instead of guessing.

  • Often the fix is data (a missing feature, mislabelled examples), not a bigger model.
  • It finds the segments where the model is worst, which averages hide.

In the real world: Reading 200 of the worst delivery-time errors, Kirana Kart's team finds most are orders to large apartment complexes, where finding the flat takes time; a 'building type' feature fixes much of it.

Tuning hyperparameters as the first response to poor accuracy. A week of tuning typically buys a little; an hour reading fifty errors often finds a missing feature worth far more.

Practise it: Improve a Low-Recall Model

Interview questions: model choice

What it is

The questions interviewers ask about model choice, from the first screening call to the senior round. Cover the answer, say yours aloud, then compare.

Why it matters

Reading an idea and explaining it are different skills.

  • Each answer below is the two or three sentences an interviewer is listening for.
  • The reasoning behind each one is in the cards above. If an answer surprises you, go back to that card.

In the real world: Interviewers rarely reward the fanciest model. They reward a model chosen from stated constraints, with a baseline and a reason to stop climbing.

Q: What model would you start with for a tabular problem?

A simple baseline, then logistic regression, then gradient-boosted trees, which are usually the best accuracy for the cost on tables. Move to neural networks only for unstructured data, sequences, huge ID spaces or very large data.

Q: When is deep learning worth it?

For images, text, audio, sequences, huge numbers of IDs needing embeddings, very large datasets, multi-task learning, or reuse of pretrained models. Otherwise trees are usually as accurate and cheaper.

Q: What is calibration and when does it matter?

Whether predicted probabilities match observed rates. It matters when the number is used directly: expected value, bidding, business thresholds, or combining models. Fix with Platt scaling or isotonic regression.

Q: Would you use an ensemble in production?

Only if the gain clearly justifies the extra training, serving, monitoring and debugging. Usually a single strong model plus rules is the better trade.

Q: What is a multi-task model?

One model with shared layers and a head per target, such as click and purchase. Rare targets learn from common ones, and the business weights the heads.

Q: How do you improve a model?

Error analysis: sample mistakes, group by cause, fix the largest group, often with data or features rather than a bigger model.

Naming a model without saying what it would replace and what it would cost.

Exercise: climb the ladder and calibrate

What it is

A short build that uses this module's ideas. Write it yourself before looking anything up; the checks at the end tell you whether it works.

Why it matters

An idea you have only read fades in a week.

  • Building it once, even in fifty lines, is what makes it yours.
  • Interviewers can tell the difference between someone who has read about a loop and someone who has debugged one.

In the real world: This comparison table, with latency next to accuracy, is what a model review at most companies asks for.

Your task. Train three rungs of the model ladder on one dataset, compare them on accuracy and cost, then calibrate the winner.

Steps

  1. Use a tabular classification dataset with at least 50,000 rows (credit default, churn, or click data).
  2. Build a rule baseline, a logistic regression and a LightGBM model; split by time if possible.
  3. Record PR-AUC, prediction latency for one row and for 1,000 rows, and training time for each.
  4. For the best model, plot a reliability table (predicted vs observed per bin) and compute expected calibration error.
  5. Calibrate with isotonic regression on a held-out set and recompute.
  6. Explain five individual predictions with SHAP.

You are done when

  • You have a table comparing three models on accuracy, latency and training time, and a stated choice.
  • Calibration error is lower after calibration, and ranking (PR-AUC) is unchanged.
  • At least one SHAP explanation surprised you, and you know why.

Stretch. Train a small neural network on the same data and add it to the comparison table.