Back

Project 09

Two-Tower Retrieval + Ranking


A two-tower retrieval model feeding a LightGBM ranking stage on MovieLens, built to keep retrieval and ranking as genuinely separate stages in the code, the way a production system would split them into independently-scaled services, even though both run in one process here.

The Approach

Retrieval and ranking stay architecturally separate on purpose, even inside one process, so the eventual split into independently-scaled services is a deployment change, not a rewrite.

The Results

A real overfitting problem, val loss bottoming out at epoch 9 then climbing back by epoch 20, was fixed with dropout and weight decay rather than papered over with best-val-loss checkpointing, and that fix made the retrieval baseline itself stronger first: Recall@10 1.4% to 3.5%, val loss now decreasing monotonically through all 20 epochs. Ranking then lifts Recall@10 further to 8.2%, a +167% relative NDCG@10 uplift from the ranking stage, on top of a retrieval baseline that had itself roughly doubled.

A ranking-invariant test asserts the ranker can only reorder retrieval's candidates, never invent a new item_id, a contract test on a non-LLM ML component in the same spirit as this portfolio's mocked-model tests elsewhere. Cold-start users, no training history, fall back to popularity-ranked items instead of crashing, verified at both the function level and through the real FastAPI endpoint. A genre-consistency sanity check, a user with 26% Sci-Fi interactions gets 6 of their top-10 recommendations tagged Sci-Fi, confirms the model learned real preference signal, not just popularity.

The Honest Parts

serving_bundle.pkl freezes features at training time, a named training-serving-skew tradeoff documented as a real production concern, not just a caveat, since a real system would refresh it on a schedule rather than once. CI trains the model fresh from scratch on every push instead of committing binary weights to git, making it a genuine regression check on the training pipeline itself, not just the serving code around it.

What I Learned

  • Fix overfitting at the source before trusting a downstream uplift numberBest-val-loss checkpointing was hiding a real overfitting problem, not solving it. Adding dropout and weight decay fixed the training curve directly, and made the retrieval baseline stronger, which is why the ranking stage's relative uplift looks smaller than an earlier run: a sign of a healthier baseline, not a weaker ranking stage, and worth being able to explain that distinction.
  • Splitting retrieval and ranking in the code pays off even before you need to split the deploymentRetrieval is memory-bound, ranking is CPU-bound. Keeping them as separate modules that share no state, even while running in one process, means the eventual move to independently-scaled services only requires a deployment change, not touching the model code.
  • A contract test can verify a ranking model without asserting exact outputAsserting the ranker only ever reorders retrieval's own candidate set, and never introduces a new item_id, catches a real class of bug (a ranking model that hallucinates or corrupts candidate identity) without needing to assert what the 'correct' ranking is.
  • Training-serving skew deserves a name, not just a footnoteCalling out that serving_bundle.pkl is a frozen snapshot, and that a real system would refresh it against a feature store, turns a limitation into a demonstrated understanding of a real production failure mode, feature drift between training and serving, rather than an implicit gap.

Tech Stack

PyTorchFAISSLightGBMFastAPIGradioW&BDocker

Links