Every consulting conversation that starts with "we have a lot of user interaction data" ends up in the same place: recommendations, or its close cousin, semantic search. Our existing tutorials cover training throughput, distribution, deployment, and generative models - but nothing covers the architecture that most revenue-facing ML systems actually use. This tutorial fills that gap: a two-tower retrieval model built with Keras 3 on TensorFlow, indexed for approximate nearest neighbour search, and served behind a real endpoint.
The pattern generalises far beyond product recommendations. Candidate retrieval for search, "related documents" panels, RAG context selection, fraud-similarity lookups, and job/candidate matching are all the same shape: learn an embedding for the query side, learn an embedding for the item side, and make the dot product mean something.
Why two towers, and where the ranker fits
A production recommender is almost always two stages:
- Retrieval. Reduce millions of items to a few hundred candidates in single-digit milliseconds. This is the two-tower model plus an ANN index. It must be cheap, so the two sides never see each other - the item tower can be precomputed offline.
- Ranking. Score those few hundred candidates with a richer model that can look at query-item interactions, business rules, and expensive features. This is where you spend your latency budget.
Teams that skip stage one and try to rank everything do not ship. Teams that skip stage two ship something that is technically correct and commercially bland. This tutorial builds stage one properly and sketches stage two.
The data
Retrieval trains on positives only: pairs of (user context, item the user engaged with). You do not need negative examples, because in-batch sampled softmax uses the other items in the batch as negatives. That property is what makes the approach practical.
import tensorflow as tf
# events.tfrecord: user_id, item_id, plus context features
def parse(record):
features = tf.io.parse_single_example(record, {
"user_id": tf.io.FixedLenFeature([], tf.string),
"item_id": tf.io.FixedLenFeature([], tf.string),
"hour_of_day": tf.io.FixedLenFeature([], tf.int64),
})
return features
ds = (
tf.data.TFRecordDataset(tf.io.gfile.glob("gs://bucket/events-*.tfrecord"))
.map(parse, num_parallel_calls=tf.data.AUTOTUNE)
.shuffle(100_000)
.batch(4096)
.prefetch(tf.data.AUTOTUNE)
)
Three data decisions matter more than any modelling choice:
- Split by time, not randomly. Train on last month, evaluate on this week. A random split leaks the future and will tell you the model is excellent right up until launch day.
- Define "engagement" honestly. A click on a carousel that autoplays is not a preference signal. Pick the strongest signal you have (purchase, completed watch, saved item) even if it is sparser.
- Batch size is a hyperparameter, not an infra setting. With in-batch negatives, a larger batch means harder negatives and a better model. 4096 is a reasonable starting point; if your batches are small, this whole approach underperforms. Our tf.data tutorial covers keeping batches that large fed.
Vocabularies and the tower bodies
Keras preprocessing layers keep the string-to-index mapping inside the model, so serving cannot drift from training:
import keras
from keras import layers
user_ids = ... # numpy array of known user id strings
item_ids = ... # numpy array of known item id strings
EMBED_DIM = 64
user_lookup = layers.StringLookup(vocabulary=user_ids, mask_token=None)
item_lookup = layers.StringLookup(vocabulary=item_ids, mask_token=None)
StringLookup maps anything unseen to index 0, so cold-start ids get a single shared "unknown" embedding instead of an exception. That is the behaviour you want in production, but it also means cold-start quality is your responsibility - lean on content features for new items.
def make_query_tower():
uid = keras.Input(shape=(), dtype="string", name="user_id")
hour = keras.Input(shape=(), dtype="int64", name="hour_of_day")
u = layers.Embedding(user_lookup.vocabulary_size(), EMBED_DIM)(user_lookup(uid))
h = layers.Embedding(24, 8)(hour)
x = layers.Concatenate()([u, h])
x = layers.Dense(128, activation="relu")(x)
x = layers.Dense(EMBED_DIM)(x)
x = layers.UnitNormalization()(x)
return keras.Model({"user_id": uid, "hour_of_day": hour}, x, name="query_tower")
def make_candidate_tower():
iid = keras.Input(shape=(), dtype="string", name="item_id")
x = layers.Embedding(item_lookup.vocabulary_size(), EMBED_DIM)(item_lookup(iid))
x = layers.Dense(128, activation="relu")(x)
x = layers.Dense(EMBED_DIM)(x)
x = layers.UnitNormalization()(x)
return keras.Model({"item_id": iid}, x, name="candidate_tower")
Two details that are easy to get wrong:
- Normalize both outputs. With
UnitNormalization, the dot product is cosine similarity, which keeps the loss well behaved and makes the ANN index's distance metric match training. - Same output dimension on both towers, obviously - but also keep it modest. 32-128 dimensions is the sweet spot; large embeddings cost index memory and query latency for very little retrieval quality.
The retrieval loss
The training objective is a softmax over the batch: for each row, the paired item should score higher than every other item in the batch.
class TwoTower(keras.Model):
def __init__(self, query_tower, candidate_tower, temperature=0.05):
super().__init__()
self.query_tower = query_tower
self.candidate_tower = candidate_tower
self.temperature = temperature
self.loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
def call(self, inputs, training=False):
q = self.query_tower(inputs, training=training)
c = self.candidate_tower(inputs, training=training)
return q, c
def train_step(self, inputs):
with tf.GradientTape() as tape:
q, c = self(inputs, training=True)
logits = tf.matmul(q, c, transpose_b=True) / self.temperature
labels = tf.range(tf.shape(logits)[0])
loss = self.loss_fn(labels, logits)
grads = tape.gradient(loss, self.trainable_variables)
self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
self.loss_tracker.update_state(loss)
return {"loss": self.loss_tracker.result()}
If you would rather not hand-roll this, the tensorflow_recommenders library packages the same objective (tfrs.tasks.Retrieval with FactorizedTopK metrics) and is still the fastest way to a baseline on the TensorFlow backend. Write the loop yourself when you need Keras 3 backend portability or a custom sampling scheme.
Popularity bias, the correction everyone forgets
In-batch negatives make popular items appear as negatives far more often than rare ones, so the model over-penalises them and under-recommends them. The standard fix is logQ correction: subtract the log sampling probability of each candidate from its logit.
log_prob = tf.gather(item_log_probs, item_indices) # precomputed from training frequencies
logits = logits - log_prob[tf.newaxis, :]
Skipping this is the single most common reason a two-tower model looks fine offline and recommends nothing but obscure long-tail items in production.
Evaluating retrieval
Accuracy is meaningless here. Use ranking metrics against a time-held-out set:
- Recall@k (k = 100, 500) - does the candidate set contain the item the user actually engaged with? This is the only metric retrieval is responsible for.
- Mean reciprocal rank for a sanity check on ordering.
- Coverage and tail share - what fraction of the catalogue ever appears in a candidate set? A model with excellent recall and 2% catalogue coverage is a popularity baseline in disguise.
Always compute the same metrics for two baselines: most-popular-items and a simple co-occurrence model. We have seen more than one client project where the neural model, once evaluated honestly, was losing to "most popular in the user's last-viewed category". Knowing that early is cheap; discovering it after launch is not.
Building the ANN index
The item tower runs offline over the catalogue, producing one vector per item:
items = tf.data.Dataset.from_tensor_slices({"item_id": item_ids}).batch(8192)
item_vectors = candidate_tower.predict(items) # (num_items, EMBED_DIM)
For catalogues up to a few hundred thousand items, exact brute-force matmul on GPU is genuinely fast enough - do that first and skip the index entirely. Beyond that, use ScaNN (via tfrs.layers.factorized_top_k.ScaNN, or the standalone scann package) for a native TensorFlow index, or an external vector store if one is already in your stack.
import scann
searcher = (
scann.scann_ops_pybind.builder(item_vectors, 100, "dot_product")
.tree(num_leaves=2000, num_leaves_to_search=100, training_sample_size=250_000)
.score_ah(2, anisotropic_quantization_threshold=0.2)
.reorder(200)
.build()
)
Tune num_leaves_to_search and reorder against measured recall-versus-latency, not defaults. Measure ANN recall against your own brute-force top-k, so you know how much retrieval quality the index is costing you - typically you can hold 95-98% for a 10-50x latency win.
Serving it
Export the query tower as its own SavedModel; it is the only half that runs per request:
query_tower.export("serving/query_tower/1")
Then the request path is: query tower forward pass (sub-millisecond on CPU for a model this size) → ANN lookup → ranker → business rules → response. Run the query tower behind TensorFlow Serving exactly as in our production serving tutorial, including dynamic batching and version labels.
Operational notes we always end up writing into the runbook:
- Rebuild the index and the query tower together. Embeddings from model version 8 are meaningless against an index built by version 7. Version them as one artifact and deploy them atomically.
- Refresh cadence. Recompute item vectors nightly for a fast-moving catalogue; retrain the towers weekly or monthly. New items need a vector before they can ever be recommended.
- Filter after retrieval, not before. Out-of-stock, region-locked, and already-seen filtering happens on the candidate list, so over-fetch (top 500, filter, keep 100) or you will return short lists.
- Log the candidate set, not just the final recommendation. When quality regresses you need to know whether retrieval or ranking broke, and no metric will tell you that after the fact.
The ranker, briefly
Once retrieval works, the ranker is comparatively simple: a model that takes query features, candidate features, and cross features, and predicts the engagement probability. Because it only sees a few hundred candidates, you can afford feature crosses, sequence features over recent history, and gradient-boosted models alongside the neural one. Train it on impressions with labels (shown-and-clicked versus shown-and-ignored) - not on the positives-only data used for retrieval. Mixing those two datasets up is the second most common failure in this architecture.
What this costs in practice
For a catalogue of a few hundred thousand items and tens of millions of events:
- Training the towers: under an hour on a single modern GPU per run, given a tf.data pipeline that can saturate it.
- Index build: minutes; it is a batch job, not a service.
- Serving: a small CPU deployment handles thousands of QPS for the query tower; the index is memory-bound, so size the host by vectors × dimensions × bytes, plus overhead.
The engineering is not the expensive part. The expensive parts are the event schema, the honest time-based evaluation harness, and the offline-online parity checks - which is exactly where we tell clients to invest first.
Where we come in
Most of the recommender work we are brought into is not "build a model" but "we built one and cannot tell whether it works". A two-week engagement to put a time-based evaluation harness, baselines, and candidate-set logging around an existing model usually changes the roadmap more than another modelling iteration would.
If you are scoping retrieval, recommendations, or semantic search on the TensorFlow stack, our TensorFlow ML consulting and services pages describe how we work - or get in touch with your data shape and traffic numbers and we will tell you what is realistic.