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¹. 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².

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².

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².

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³. 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².

flowchart LR
    User[User request] --> UF[User feature fetch]
    UF --> R1[Two-tower ANN]
    UF --> R2[User history sims]
    UF --> R3[Trending heuristics]
    R1 --> MX[Candidate mixer]
    R2 --> MX
    R3 --> MX
    MX -->|~1000 candidates| S1[Stage 1 ranker]
    S1 -->|top 200| S2[Stage 2 ranker: MTML]
    S2 -->|top 50| RR[Re-ranking: diversity + policy]
    RR -->|final page| Render[User sees results]

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⁴.

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⁴. 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.

flowchart TD
    subgraph Training
        UI[User features] --> UT[User tower MLP]
        II[Item features] --> IT[Item tower MLP]
        UT --> DP[Dot product]
        IT --> DP
        DP --> LOSS[In-batch softmax + log-Q correction]
    end
    subgraph Serving
        UT2[User tower: online] --> UV[User vector]
        IT2[Item tower: offline batch] --> IV[(ANN index: FAISS / ScaNN)]
        UV --> ANN[Top-K query]
        IV --> ANN
        ANN --> CANDS[~1000 candidates to ranker]
    end

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⁵.

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)⁶.

Feature crosses come in two flavors. Explicit crosses are hand-crafted (country x item_category), memorized by a wide linear layer in Wide & Deep⁷. 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⁸; DLRM builds second-order dot products between every pair of categorical embedding vectors⁹.

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.

flowchart LR
    RAW[(Raw events: Kafka)] --> T[Shared transform library]
    T --> OFF[(Offline store: warehouse)]
    T --> ONL[(Online store: Redis / DynamoDB)]
    OFF -->|point-in-time join| TRAIN[Training pipeline]
    ONL -->|low-latency lookup| SERVE[Online inference]
    SERVE -->|shadow log| MON[Skew monitor: PSI / KS]
    OFF --> MON
    MON -->|alert on drift| ONCALL[On-call]

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¹.

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¹⁰.

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.

stateDiagram-v2
    [*] --> Registry: publish artifact
    Registry --> Shadow: deploy, log predictions only
    Shadow --> Canary1: offline-online parity check passes
    Canary1 --> Canary5: 1% to 5% traffic
    Canary5 --> Canary25: guardrails green
    Canary25 --> Full: 25% to 100% traffic
    Canary1 --> Rollback: guardrail regression
    Canary5 --> Rollback
    Canary25 --> Rollback
    Rollback --> Registry
    Full --> [*]

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¹¹; 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¹². 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³.

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³.

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³.

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)², Pinterest PinSage (3 billion nodes, 18 billion edges)¹³, Meta DLRM (TB-scale sparse embeddings, thousands of GPUs)¹⁴, and TikTok Monolith (minute-level online training with collisionless cuckoo hashing)⁵.

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)¹¹, and Netflix subsequently reported interleaving needing more than 100x fewer subscribers than their most sensitive A/B metric for ranking comparisons¹². 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