ML System Design Fundamentals

TL;DR: Every large-scale recommendation feed runs a two-stage pipeline: retrieval narrows billions of items to thousands in under 50 ms, then a ranking model scores those thousands with rich features in about 100 ms. Two-tower embedding models power retrieval by precomputing item vectors and querying an ANN index at request time. The feature store keeps training and serving consistent. Training-serving skew, not the model itself, is the #1 silent killer of production ML: Sculley et al. found that only a small fraction of real-world ML code is the model; the rest is surrounding infrastructure that can silently diverge^[1]. Get the pipeline right before you tune the model.

Learning Objectives#

After this module, you will be able to:

• Explain why the candidate-generation then ranking two-stage pattern outperforms a single model over the full catalog
• Describe a two-tower model: what each tower computes, how training works, how serving does dot-product retrieval
• Design a feature pipeline handling categorical, numeric, embedding, and sequence features at scale
• Identify training-serving skew and pick a mitigation (feature store, shared transform, parity tests)
• Plan a model rollout using versioning, shadow deploys, and a staged A/B test
• Reason about interleaving vs classic A/B testing and when each applies

Intuition#

Imagine you run a bookstore with 100 million titles. A customer walks in and says, "I want something good to read tonight." You cannot hand them every book and ask them to rank all 100 million. Instead, you do two things. First, a fast librarian scans the shelves and pulls 500 books that roughly match the customer's taste based on their reading history and the book's genre tags. This takes seconds. Second, a slower but more discerning expert examines those 500 books, considering the customer's mood, the time of day, recent purchases, and subtle cross-signals ("people who liked X also liked Y"), then arranges the top 20 on a display table.

The fast librarian is the retrieval stage. She optimizes for recall: do not miss the right book. The expert is the ranking stage. He optimizes for precision: get the order right on a small set. Neither can do the other's job. The librarian cannot spend 30 seconds per book (that is 95 years for 100 million titles). The expert cannot evaluate 100 million books in real time.

This is the fundamental architecture behind YouTube's home feed, Instagram Explore, TikTok's For You page, Pinterest's home feed, and every ads auction at scale. The rest of this chapter explains how to build it: the retrieval model (two-tower embeddings), the ranking model (DLRM, DCN, attention), the feature pipeline that feeds both, and the infrastructure that keeps them honest.

Theory#

The two-stage retrieval-then-ranking pattern#

A single deep model cannot score 10^9 items per request. At 1 microsecond per forward pass, that is 1,000 seconds per request. The two-stage pattern exists because latency budgets and model cost multiply.

Stage 1 (retrieval) narrows the full catalog to roughly 1,000 candidates in under 50 ms. It uses cheap, recall-oriented scorers: a two-tower ANN lookup, an inverted index over user interaction history, or trending heuristics. Multiple retrieval sources run in parallel and a candidate mixer deduplicates their outputs^[2].

Stage 2 (ranking) scores those 1,000 candidates with a precision-oriented deep model that can afford rich cross features (user-item interactions, context, real-time signals). Budget: about 100 ms^[2:1].

Stage 3 (re-ranking) applies diversity constraints, integrity filters (harmful content, policy), and business rules. Instagram Explore enforces "no two posts from the same author in a row" here^[2:2].

YouTube's 2016 paper made this split canonical: candidate generation was a deep MLP producing a 256-dimensional user embedding with softmax over millions of videos; ranking was a separate network scoring candidates using impression features and watch-time-weighted logistic loss^[3]. Instagram Explore (2023) formalized four sub-stages: retrieval, lightweight two-tower first-pass ranker (distilled from stage 2), heavy multi-task multi-label (MTML) second-pass ranker, and final re-ranking^[2:3].

The funnel: billions of items become thousands via retrieval, hundreds via ranking, and a final page via re-ranking with diversity and policy filters.

Two-tower embedding models#

A two-tower model has two independent subnetworks. The user tower encodes user context (watch history, demographics, device) into a d-dimensional vector (typically 64 to 256). The item tower encodes item features (title, category, popularity) into the same space. They join only at the top via dot product or cosine similarity^[4].

Training: Each minibatch contains (user, positive item) pairs. Other items in the batch act as in-batch negatives. The loss is softmax over the batch. Yi et al. (2019) showed that without correction, popular items appear as negatives in many batches and get over-penalized, producing a biased estimate of the full-softmax gradient. Their fix: subtract log(P_sampled(item)) from logits (log-Q correction) using a streaming frequency estimator^[4:1]. Hard-negative mining replaces random negatives with items the model currently scores highly but the user did not engage with.

Serving: The item tower runs offline over the full catalog, producing one vector per item. Those vectors load into an ANN index (FAISS, ScaNN, or HNSW, as covered in Vector Databases). At request time, only the user tower runs, produces one vector, and the ANN index returns top-K in sub-linear time.

User and item towers train jointly on in-batch softmax; at serving time the item tower runs offline and the user tower runs online against a precomputed ANN index.

The key constraint: because the towers are independent, a two-tower model cannot see interaction features ("this user watched a similar video 20 seconds ago"). Those features belong in the ranker.

Feature engineering at scale#

Production ranking models consume four families of features:

Categorical (user_id, country, item_category): lookup into an embedding table of size vocab_size x d. For high-cardinality IDs (100M users, 10^10 ad creatives), you either hash into a smaller table (the hashing trick, with collisions) or use a collisionless cuckoo-hash table as TikTok's Monolith does^[5].

Numeric (price, age, session length): log transform log(1 + x) for heavy-tailed distributions, quantile normalization, or equal-frequency bucketization feeding an embedding.

Pre-trained embeddings (content, text): CLIP vectors for images, Sentence-BERT for text. Frozen or fine-tuned, concatenated with learned embeddings.

Sequence (last N watches, last N queries): mean-pool, self-attention (transformer encoder as in BST and SASRec), or target attention over the candidate (DIN from Alibaba)^[6].

Feature crosses come in two flavors. Explicit crosses are hand-crafted (country x item_category), memorized by a wide linear layer in Wide & Deep^[7]. Learned crosses are automatic: DCN-V2 uses mixture-of-low-rank cross layers delivering offline AUC and online business-metric lifts across multiple Google systems^[8]; DLRM builds second-order dot products between every pair of categorical embedding vectors^[9].

Offline and online feature stores#

A feature store materializes and serves features twice: once for training (offline, historical, point-in-time correct) and once for inference (online, low-latency, freshest-version).

The offline store backs onto a warehouse (Snowflake, BigQuery, Iceberg). The online store backs onto a KV system (Redis, DynamoDB, ScyllaDB). Both must return the same value at the same logical time. That is the contract.

The critical property is point-in-time correctness: to train a model predicting whether user U clicks item I at time T, you must join to U's features as of time T minus one second, never T plus one hour. Without this, training AUC is inflated by 5 to 10 points that vanish in A/B because you leaked future information into training rows.

A shared transform library writes to both offline and online stores; a skew monitor continuously diffs the distributions to catch divergence before it reaches production.

Feast is the reference open-source feature store; managed options include Tecton (now part of Databricks), Hopsworks, Databricks Feature Store, Vertex AI Feature Store (Bigtable online serving), and SageMaker Feature Store. Each feature has a freshness SLO: "user_last_click_item_id must be less than 5 seconds old" for a ranking feature vs "user_signup_country is static."

Training-serving skew#

Training-serving skew is a divergence between the feature distribution the model trained on and what it sees at inference. It is caused by two code paths computing the "same" feature differently. Sculley et al. (2015) flagged this as the signature pathology of applied ML^[1:1].

Root causes in descending order of frequency:

1. Separate transform codebases (Spark UDF in training, Python lambda at serving)
2. Unit mismatches (miles vs kilometers, seconds vs milliseconds)
3. Null handling (train fillna(0), serve fillna(-1))
4. Timezone drift (training in UTC, serving in local)
5. Vocabulary drift (new category at serving absent from training)

Detection: Shadow-log online features back to the offline store, then diff distributions with Population Stability Index (PSI), Kolmogorov-Smirnov (KS), or TFDV's DetectFeatureSkew^[10].

Prevention: Shared transform library (single source of truth), online features fed back into training via the feature store, and contract tests asserting offline_transform(x) == online_transform(x) on a held-out sample before every deploy.

The post-mortem pattern is always the same: "we shipped, CTR dropped 2%, we spent two weeks debugging, and discovered user_age was in years offline and months online."

Model rollout and A/B testing#

Each model artifact includes weights, feature schema, training data snapshot hash, offline metrics, and serving config. The rollout sequence:

1. Shadow deploy runs the new model on 100% of traffic but does not serve its predictions; predictions are logged for offline comparison.
2. Canary gates on metrics: 1% to 5% to 25% to 50% to 100% over hours to days, pausing if any guardrail (latency p99, error rate, fairness slice, business metric) regresses.
3. Rollback is a feature flag flip, not a redeploy.

A new model version moves through four gates with automated rollback on any guardrail regression. Rollback is always a feature-flag flip back to the prior artifact in the registry, never a redeploy.

For online evaluation, classic A/B randomizes users into treatment and control and measures a primary metric (CTR, retention, revenue). Interleaving (Chapelle, Joachims, Radlinski, Yue 2012^[11]; Netflix 2017) blends two rankers' outputs on the same page and attributes each click to whichever ranker contributed that position. Netflix reported that interleaving requires more than 100x fewer subscribers than their most sensitive A/B metric to detect a ranking difference^[12]. Use interleaving as a pre-filter for ranking changes; use classic A/B for architectural or marketplace-affecting launches.

Interference breaks classic A/B when treatment users interact with control users (social networks) or treatment bids affect control prices (ads auctions). Mitigations: cluster-randomization (whole social clusters get the same treatment) or switchback designs (flip regions on/off in time slices).

Real-World Example#

YouTube's two-stage recommender (Covington, Adams, Sargin 2016) remains the canonical reference. At the time of publication, YouTube had hundreds of millions of videos and billions of user-video interactions per day^[3:1].

The candidate-generation network takes the user's watch history (bag of video embeddings, averaged), search tokens, demographics, and an "example age" feature into several fully-connected ReLU layers (the paper sweeps depth; one of the configurations reported in Section 3.5 widens to 1024, then 512, then 256 units) that output a 256-dimensional user embedding. During training, softmax runs over the video vocabulary. At serving time, precomputed video vectors are loaded into an ANN index, and the user embedding retrieves the top few hundred candidates^[3:2].

The ranking network consumes hundreds of features per (user, video) pair, including impression video ID, user language, device, and time since last watch. Critically, the ranking objective is expected watch time (regression with weighted logistic loss), not click probability, because watch time correlated with business value better than CTR^[3:3].

Two engineering decisions stand out. First, the "example age" feature: a continuous value set to the time between video upload and training time. At inference it is set to zero to represent "now," which materially improved recommendation freshness. Second, training labels are the next watched video with negatives sampled from the background video distribution and importance-weighted, which reduces to nearest-neighbor search at serving time.

This architecture, published in 2016, spawned every subsequent two-stage system: Instagram Explore (four-stage retrieve-rank-rerank pipeline at Meta scale, with two-tower distillation in stage 2 and multi-task learning in stage 3)^[2:4], Pinterest PinSage (3 billion nodes, 18 billion edges)^[13], Meta DLRM (TB-scale sparse embeddings, thousands of GPUs)^[14], and TikTok Monolith (minute-level online training with collisionless cuckoo hashing)^[5:1].

Trade-offs#

Interleaving is an evaluation methodology rather than a ranking architecture, so it belongs with the A/B discussion: each user sees results from both arms on the same query, removing between-user variance and dramatically increasing sensitivity. Chapelle et al. report that one expert-judged query is roughly equivalent to ten clicked queries under interleaving (~10x more sensitive than judged-query offline evaluation)^[11:1], and Netflix subsequently reported interleaving needing more than 100x fewer subscribers than their most sensitive A/B metric for ranking comparisons^[12:1]. It provides no retention or long-term-outcome signal and is only applicable to ranking changes, so run it as a pre-filter on candidate ranking models before committing to a full A/B.

Common Pitfalls#

Exercise#

Design a news-ranking service for 100 million DAU over ~10 million articles from the last 30 days. Specify: (a) the two-tower model, with loss and negative sampling; (b) the ranker feature set split by categorical, numeric, embedding, and sequence; (c) where each feature lives in the feature store and its freshness SLO; (d) how you detect training-serving skew; (e) the rollout plan with primary and guardrail metrics.

Key Takeaways#

• Two-stage retrieval-then-ranking is the dominant architecture for large-catalog personalization. It exists because scoring 10^9 items with a deep model per request is physically impossible within latency budgets.
• Two-tower models with dot-product retrieval are the workhorse of candidate generation. Get the negative sampling and log-Q correction right, and the rest falls into place.
• The feature store is the contract that keeps training and serving honest. Without it, skew is inevitable.
• Training-serving skew is almost always a process problem (two codebases, two data paths) before it is a modeling problem. Check the pipeline before you retrain the model.
• Interleaving is roughly 100x more sensitive than classic A/B for ranking changes, but it cannot measure retention or marketplace effects.
• Point-in-time correctness in training data is non-negotiable. Leaking future features inflates offline metrics by 5 to 10 points that vanish online.
• Model rollout is a feature flag problem: shadow, canary, ramp, rollback. Never ship without a one-flag revert path.

Further Reading#

• Deep Neural Networks for YouTube Recommendations (Covington et al. 2016). the foundational two-stage paper; every later system references its candidate-generation and ranking split.
• Sampling-Bias-Corrected Neural Modeling for Large Corpus Item Recommendations (Yi et al. 2019). canonical two-tower paper with log-Q correction; read this to understand why naive in-batch softmax fails.
• Hidden Technical Debt in Machine Learning Systems (Sculley et al. 2015). names training-serving skew and the "95% glue code" observation; essential framing for feature stores and monitoring.
• Scaling the Instagram Explore Recommendations System (Meta 2023). four-stage ranker with two-tower distillation; the most modern public reference for production multi-stage ranking.
• Monolith: Real Time Recommendation System with Collisionless Embedding Table (Liu et al. 2022). ByteDance's online-training system with cuckoo hashing; shows how the "retrain daily" assumption breaks at TikTok pace.
• Feast Documentation. reference open-source feature store; covers offline/online split, point-in-time correctness, and materialization patterns.
• Innovating Faster on Personalization Algorithms at Netflix Using Interleaving. Netflix's evidence that interleaving needs 100x fewer users than classic A/B for ranking evaluation.
• Designing Machine Learning Systems (Chip Huyen, O'Reilly 2022). book-length treatment of feature stores, skew, deployment, and the full ML lifecycle.

Flashcards#

References#