Most teams that build a recommender stop after retrieval. They train a two-tower model, index the embeddings, pull back the top 500 candidates in a few milliseconds, and ship it. Then engagement plateaus, and nobody can explain why a clearly better item sits at position 9.
The answer is usually that there is no ranking stage. Retrieval is optimised for recall over millions of items with a cheap dot product; ranking is optimised for order over a few hundred candidates, with expensive cross features the retrieval model cannot afford. This tutorial builds that second stage on TensorFlow, and the part that surprises people is that it is not a classification problem.
If you have not built the first stage yet, start with Two-Tower Retrieval and come back.
Why pointwise training underperforms
The instinct is to train a binary classifier: features in, P(click) out, sort by probability. It works, and it is nearly always worse than it needs to be, for two reasons.
- The loss does not know about the list. Binary cross-entropy penalises a badly scored item at position 300 exactly as hard as one at position 1. Users only see the top handful. The metric you care about (NDCG@10, MRR, top-1 precision) is dominated by the head of the list; the loss should be too.
- The label is relative, not absolute. "Item A is better than item B for this query" is a much cleaner signal than "item A is good". Absolute click rates drift with traffic mix, seasonality, and UI changes; within-list preferences are far more stable.
Listwise losses fix both. They take a whole list of candidates for one query, produce one score per candidate, and compute a loss over the induced permutation.
Shaping the data: groups, not rows
This is the step that breaks most implementations. A ranking example is not a row, it is a list: one query context plus N candidates plus N labels. Your tf.data pipeline must preserve that grouping.
import tensorflow as tf
LIST_SIZE = 50
feature_spec = {
# context: one value per list, broadcast across candidates
"query_embedding": tf.io.FixedLenFeature([64], tf.float32),
"user_segment": tf.io.FixedLenFeature([], tf.int64),
# per-candidate: LIST_SIZE values each
"item_id": tf.io.FixedLenFeature([LIST_SIZE], tf.int64),
"item_price": tf.io.FixedLenFeature([LIST_SIZE], tf.float32),
"bm25_score": tf.io.FixedLenFeature([LIST_SIZE], tf.float32),
"retrieval_rank": tf.io.FixedLenFeature([LIST_SIZE], tf.int64),
"label": tf.io.FixedLenFeature([LIST_SIZE], tf.float32), # 0/1/2 grades
"mask": tf.io.FixedLenFeature([LIST_SIZE], tf.float32), # 1 = real, 0 = padding
}
def parse(record):
ex = tf.io.parse_single_example(record, feature_spec)
labels = ex.pop("label")
return ex, labels
ds = (tf.data.TFRecordDataset(files, num_parallel_reads=tf.data.AUTOTUNE)
.map(parse, num_parallel_calls=tf.data.AUTOTUNE)
.shuffle(10_000)
.batch(64)
.prefetch(tf.data.AUTOTUNE))
Three things to get right:
- Fixed list size with padding. Pad short lists to
LIST_SIZEand carry an explicit mask. Never let padded slots contribute to the loss. - Sample the negatives deliberately. Do not put all 500 retrieval candidates in the list. Take every positive plus a sample of the candidates that were actually shown alongside them. Items the user never saw are not confirmed negatives.
- Graded labels beat binary ones. If you have impression / click / add-to-cart / purchase, encode them as 0/1/2/3 rather than collapsing to clicked-or-not. NDCG can use the grades; binary loss throws the information away.
See tf.data Pipelines That Keep Your GPU Busy for the throughput side; grouped records are larger and I/O-bound more often than you expect.
A scoring model in Keras 3
The model scores each candidate independently but is trained over the list. Broadcast the context features across the candidate axis and run a shared MLP.
import keras
from keras import layers, ops
def build_ranker(num_items, list_size=LIST_SIZE, emb_dim=32):
query_emb = keras.Input(shape=(64,), name="query_embedding")
segment = keras.Input(shape=(), dtype="int64", name="user_segment")
item_id = keras.Input(shape=(list_size,), dtype="int64", name="item_id")
price = keras.Input(shape=(list_size,), name="item_price")
bm25 = keras.Input(shape=(list_size,), name="bm25_score")
item_vec = layers.Embedding(num_items, emb_dim)(item_id) # (B, L, E)
seg_vec = layers.Embedding(64, 8)(segment) # (B, 8)
ctx = layers.Concatenate()([query_emb, seg_vec]) # (B, 72)
ctx = ops.repeat(ops.expand_dims(ctx, 1), list_size, axis=1) # (B, L, 72)
dense_feats = ops.stack([price, bm25], axis=-1) # (B, L, 2)
x = layers.Concatenate(axis=-1)([ctx, item_vec, dense_feats])
for units in (256, 128, 64):
x = layers.Dense(units, activation="relu")(x)
x = layers.Dropout(0.1)(x)
scores = layers.Dense(1)(x) # (B, L, 1)
scores = ops.squeeze(scores, axis=-1) # (B, L)
return keras.Model(
inputs={"query_embedding": query_emb, "user_segment": segment,
"item_id": item_id, "item_price": price, "bm25_score": bm25},
outputs=scores,
)
Note the output shape: (batch, list_size), one score per candidate. No sigmoid. Ranking scores are unbounded and only their order matters — calibrating them into probabilities is a separate job (and only needed if you multiply the score by a bid or a margin downstream).
Listwise loss
The simplest listwise loss is a softmax cross-entropy over the list, sometimes called ListNet or the softmax loss. It is a strong default and is easy to write yourself, which matters because tensorflow_ranking has historically lagged behind current TensorFlow releases — check its compatibility against your pinned version before you take the dependency.
NEG_INF = -1e9
class ListwiseSoftmaxLoss(keras.losses.Loss):
"""Softmax cross-entropy over each candidate list, padding-aware."""
def call(self, y_true, y_pred):
mask = ops.cast(y_true >= 0.0, "float32") # use -1 as the pad label
scores = ops.where(mask > 0, y_pred, NEG_INF)
labels = ops.maximum(y_true, 0.0) * mask
# normalise graded labels into a target distribution per list
denom = ops.sum(labels, axis=-1, keepdims=True)
target = labels / ops.maximum(denom, 1e-9)
log_probs = scores - ops.logsumexp(scores, axis=-1, keepdims=True)
per_list = -ops.sum(target * log_probs, axis=-1)
has_pos = ops.cast(ops.squeeze(denom, -1) > 0.0, "float32")
return ops.sum(per_list * has_pos) / ops.maximum(ops.sum(has_pos), 1.0)
Lists with no positive at all contribute nothing — keep some of them in the data for realism, but they carry no gradient here.
When the softmax loss plateaus, the next step up is a metric-aligned surrogate such as ApproxNDCG, which replaces the non-differentiable rank with a smooth approximation of it. It usually buys another point or two of NDCG on head-heavy traffic, at the cost of being fiddlier to tune (the temperature matters). Pairwise logistic loss is the third option and is the most robust when labels are sparse and noisy.
Position bias, or your model learns your old UI
Clicks are biased by position: item 1 gets clicked because it was item 1. Train naively on logged clicks and the ranker learns to reproduce whatever ranked the training data — including its mistakes.
The cheap, effective fix is a two-tower debiasing setup during training only: a tiny tower sees only position-like features (retrieval_rank, device, slot), the main tower sees everything else, and the training logit is their sum. At export time you drop the bias tower and serve the relevance tower alone.
bias_in = keras.Input(shape=(LIST_SIZE,), dtype="int64", name="retrieval_rank")
bias_logit = ops.squeeze(layers.Embedding(200, 1)(bias_in), axis=-1) # (B, L)
train_model = keras.Model(
inputs={**relevance_model.input, "retrieval_rank": bias_in},
outputs=relevance_model.output + bias_logit,
)
train_model.compile(optimizer=keras.optimizers.Adam(1e-3), loss=ListwiseSoftmaxLoss())
train_model.fit(ds, validation_data=val_ds, epochs=10)
# serve relevance_model, not train_model
If you can afford it, a small slice of randomised or swapped-position traffic gives you an unbiased evaluation set and a direct estimate of the propensity curve. It is the highest-value 1% of traffic you will ever give away.
Evaluating: NDCG, not AUC
AUC over pooled rows tells you almost nothing about list quality. Compute per-list metrics and average them.
def ndcg_at_k(y_true, y_pred, k=10):
mask = ops.cast(y_true >= 0.0, "float32")
labels = ops.maximum(y_true, 0.0) * mask
scores = ops.where(mask > 0, y_pred, NEG_INF)
order = ops.argsort(-scores, axis=-1)[:, :k]
gains = ops.take_along_axis(labels, order, axis=-1)
discounts = 1.0 / ops.log2(ops.arange(2, k + 2, dtype="float32"))
dcg = ops.sum((ops.power(2.0, gains) - 1.0) * discounts, axis=-1)
ideal = ops.sort(labels, axis=-1)[:, ::-1][:, :k]
idcg = ops.sum((ops.power(2.0, ideal) - 1.0) * discounts, axis=-1)
return ops.mean(dcg / ops.maximum(idcg, 1e-9))
Report NDCG@10 and MRR against three baselines: the current production ranker, the raw retrieval order, and a random shuffle. The retrieval-order baseline is the one that keeps you honest — if your ranker cannot beat "leave the candidates in the order retrieval produced", it is not earning its latency.
Slice the metrics: new users vs. returning, head vs. tail queries, short vs. long lists. An overall gain built entirely on head queries is a fragile result, and slicing is also how you catch the fairness problems that a single aggregate number hides.
Exporting for serving
Serve the ranker with a signature that scores a whole list in one call. Per-candidate calls will destroy your latency budget.
@tf.function(input_signature=[{
"query_embedding": tf.TensorSpec([None, 64], tf.float32),
"user_segment": tf.TensorSpec([None], tf.int64),
"item_id": tf.TensorSpec([None, None], tf.int64),
"item_price": tf.TensorSpec([None, None], tf.float32),
"bm25_score": tf.TensorSpec([None, None], tf.float32),
}])
def rank(features):
scores = relevance_model(features, training=False)
order = tf.argsort(scores, axis=-1, direction="DESCENDING")
return {"scores": scores, "order": order,
"ranked_item_id": tf.gather(features["item_id"], order, batch_dims=1)}
tf.saved_model.save(relevance_model, "serving/ranker/1",
signatures={"serving_default": rank})
Note the dynamic candidate axis — production lists vary in length and you do not want a re-trace per size. From here it is the standard deployment path: dynamic batching, version labels, and canaries, all covered in Serving TensorFlow Models in Production. A ranking MLP over a few hundred candidates is small enough to run comfortably on CPU; check before you pay for GPU inference.
Keeping it alive
Rankers decay faster than almost any other model, because the moment you deploy one it changes the data that trains its successor. Three habits keep that under control:
- Retrain on a schedule, not on a vibe. Wire it into a TFX pipeline with a blessing step so no unevaluated ranker reaches production — see Continuous Training with TFX.
- Monitor the input features, not just the output. Score distributions shift late; feature distributions shift first. TensorFlow Data Validation on the candidate features catches a broken upstream join days before NDCG moves.
- Log the candidate set, not just the click. If you only log what was clicked, you cannot ever rebuild a training list. This is the single most common reason a ranking project stalls for a quarter.
Where this goes wrong in practice
The failure we are called in to fix is almost never the model architecture. It is logging: no candidate sets, no impressions, no position data, so there is nothing to train a listwise model on and the first six weeks of the project are instrumentation. If you are planning a ranking stage, fix the logging one release before you start modelling.
If you want help scoping that — the logging schema, the offline evaluation harness, or the ranker itself — our TensorFlow ML consultants have built these for search, marketplace, and feed products. Get in touch with a description of your candidate volume and latency budget and we will tell you whether a ranking stage is worth it before we quote for one.