Two-Tower Recommendation Systems

The Story That Explains Two-Tower Systems

Before Two-Tower models existed, recommenders relied on collaborative filtering, matrix factorisation, and hand-engineered feature pipelines. These worked well at small scale but collapsed under the weight of billions of users and items. The Two-Tower architecture solved this by separating the encoding of users and items — making retrieval blazing fast at inference time.

Architecture — The Two-Tower Blueprint

A Two-Tower model consists of two separate neural networks (towers) that are trained jointly but run independently at inference time. Here is the full architecture at a glance:

Candidate Generation — The Retrieval Funnel

In a full production pipeline, candidate generation is typically a multi-source funnel. Multiple candidate generators run in parallel, their outputs are de-duplicated, and then the combined pool is ranked.

The Three Goals of Candidate Generation

Retrieval Systems — From Query to Candidates

Once the Query Tower produces a user vector, we need to find the nearest item vectors. This is the retrieval problem. The naive approach (brute-force dot product over all items) is exact but prohibitively slow at scale. Production systems use Approximate Nearest Neighbour (ANN) algorithms that trade a tiny amount of accuracy for enormous speed gains.

The Major ANN Index Types

How HNSW Works Visually

Vector Similarity Search — The Mathematics

At the heart of Two-Tower retrieval is a single question: how similar are two vectors? Three metrics dominate in practice. Understanding when to use each is critical.

Why Embeddings Work — The Geometry of Preference

Training — Loss Functions and Sampling Strategy

Training a Two-Tower model is fundamentally a metric learning problem: push user and item embeddings for positive (user-item) pairs together, and push negatives apart. The choice of loss function and negative sampling strategy dramatically affects model quality.

Project — Building a YouTube / Spotify-Style Recommender

We'll build a complete Two-Tower recommender system using TensorFlow Recommenders (TFRS) and FAISS. The architecture mirrors what YouTube and Spotify use in production, adapted for the MovieLens 1M dataset. By the end, you will have a working retrieval system with an ANN index.

Step 1 — Installation and Dataset

# Install required libraries
pip install tensorflow-recommenders tensorflow-datasets faiss-cpu

import tensorflow as tf
import tensorflow_recommenders as tfrs
import tensorflow_datasets as tfds
import numpy as np
import faiss
from typing import Dict, Text

# Load MovieLens 1M — ratings + movie metadata
ratings = tfds.load("movielens/1m-ratings", split="train")
movies  = tfds.load("movielens/1m-movies", split="train")

# Extract only what we need
ratings = ratings.map(lambda x: {
    "user_id":    x["user_id"],
    "movie_id":   x["movie_id"],
    "movie_title": x["movie_title"],
    "genres":     x["genres"],
})

movies = movies.map(lambda x: {
    "movie_id":    x["movie_id"],
    "movie_title": x["movie_title"],
    "genres":      x["genres"],
})

# Build vocabularies
user_ids    = ratings.map(lambda x: x["user_id"])
movie_ids   = ratings.map(lambda x: x["movie_id"])
movie_titles = movies.map(lambda x: x["movie_title"])

user_ids_vocab    = tf.keras.layers.StringLookup(mask_token=None)
user_ids_vocab.adapt(user_ids)

movie_ids_vocab   = tf.keras.layers.StringLookup(mask_token=None)
movie_ids_vocab.adapt(movie_ids)

movie_titles_vocab = tf.keras.layers.StringLookup(mask_token=None)
movie_titles_vocab.adapt(movie_titles)

print(f"Users: {user_ids_vocab.vocabulary_size()}")
print(f"Movies: {movie_ids_vocab.vocabulary_size()}")

Step 2 — Building the Query Tower

class QueryTower(tf.keras.Model):
    """
    Encodes user context into a dense embedding vector.
    Mirrors YouTube-style query encoding with:
      - User ID embedding (identity signal)
      - Averaged watch history embedding (sequence signal)
      - Dense projection layers
    """
    def __init__(self, user_vocab_size, embedding_dim=64, output_dim=128):
        super().__init__()
        self.user_embedding = tf.keras.Sequential([
            tf.keras.layers.StringLookup(
                vocabulary=user_ids_vocab.get_vocabulary(), mask_token=None
            ),
            tf.keras.layers.Embedding(user_vocab_size, embedding_dim),
        ])
        # Dense layers for projection
        self.dense_layers = tf.keras.Sequential([
            tf.keras.layers.Dense(256, activation="relu"),
            tf.keras.layers.Dense(128, activation="relu"),
            tf.keras.layers.Dense(output_dim),  # No activation — L2 norm next
        ])

    def call(self, inputs):
        user_emb = self.user_embedding(inputs["user_id"])
        x = self.dense_layers(user_emb)
        # L2 normalise → dot product = cosine similarity
        return tf.nn.l2_normalize(x, axis=1)

Step 3 — Building the Item Tower

class ItemTower(tf.keras.Model):
    """
    Encodes items (movies) into a dense embedding vector.
    Uses both:
      - Movie ID embedding (collaborative signal)
      - Title text embedding via bag-of-words (content signal)
      - Genre embedding (metadata signal)
    Pre-computed offline and stored in ANN index.
    """
    def __init__(self, movie_vocab_size, title_vocab_size,
                 embedding_dim=64, output_dim=128):
        super().__init__()
        # Movie ID embedding (ID-based collaborative signal)
        self.movie_embedding = tf.keras.Sequential([
            tf.keras.layers.StringLookup(
                vocabulary=movie_ids_vocab.get_vocabulary(), mask_token=None
            ),
            tf.keras.layers.Embedding(movie_vocab_size, embedding_dim),
        ])
        # Title text — bag-of-words over subword tokens
        self.title_embedding = tf.keras.Sequential([
            tf.keras.layers.StringLookup(
                vocabulary=movie_titles_vocab.get_vocabulary(), mask_token=None
            ),
            tf.keras.layers.Embedding(title_vocab_size, embedding_dim),
            tf.keras.layers.GlobalAveragePooling1D(),  # mean over token embeddings
        ])
        self.dense_layers = tf.keras.Sequential([
            tf.keras.layers.Dense(256, activation="relu"),
            tf.keras.layers.Dense(128, activation="relu"),
            tf.keras.layers.Dense(output_dim),
        ])

    def call(self, inputs):
        movie_emb = self.movie_embedding(inputs["movie_id"])
        title_emb = self.title_embedding(inputs["movie_title"])
        # Concatenate ID + content signals
        combined = tf.concat([movie_emb, title_emb], axis=1)
        x = self.dense_layers(combined)
        return tf.nn.l2_normalize(x, axis=1)

Step 4 & 5 — Assembling and Training the Full Model

class TwoTowerModel(tfrs.Model):
    """Full Two-Tower retrieval model with TFRS retrieval task."""

    def __init__(self, query_tower, item_tower, movies_dataset):
        super().__init__()
        self.query_tower = query_tower
        self.item_tower  = item_tower

        # Retrieval task: in-batch softmax + factorised top-K metrics
        self.task = tfrs.tasks.Retrieval(
            metrics=tfrs.metrics.FactorizedTopK(
                candidates=movies_dataset.batch(128).map(self.item_tower)
            )
        )

    def compute_loss(self, features, training=False):
        user_embeddings  = self.query_tower(features)
        movie_embeddings = self.item_tower(features)
        return self.task(user_embeddings, movie_embeddings,
                        compute_metrics=not training)

# Instantiate towers
query_tower = QueryTower(user_ids_vocab.vocabulary_size())
item_tower  = ItemTower(
    movie_ids_vocab.vocabulary_size(),
    movie_titles_vocab.vocabulary_size()
)

# Build full model
model = TwoTowerModel(query_tower, item_tower, movies)
model.compile(optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.1))

# Train / validation split
tf.random.set_seed(42)
shuffled = ratings.shuffle(100_000, seed=42, reshuffle_each_iteration=False)
train = shuffled.take(80_000)
val   = shuffled.skip(80_000).take(20_000)

# Train — 10 epochs, batch size 4096
history = model.fit(
    train.batch(4096),
    validation_data=val.batch(4096),
    epochs=10,
    verbose=1,
)

Step 6 — Building the FAISS ANN Index

# Step 6a: Generate embeddings for ALL movies (Item Tower inference)
movie_ids_all    = []
movie_titles_all = []

for movie in movies.batch(512):
    movie_ids_all.extend(movie["movie_id"].numpy())
    movie_titles_all.extend(movie["movie_title"].numpy())

# Run Item Tower over all movies
movie_embeddings = item_tower.predict(
    tf.data.Dataset.from_tensor_slices({
        "movie_id":    movie_ids_all,
        "movie_title": movie_titles_all,
    }).batch(512)
)  # shape: (3884, 128) — one 128-d vector per movie

# Step 6b: Build HNSW index with FAISS
DIM = 128         # embedding dimension
M   = 32          # HNSW graph connectivity (higher = better recall, more memory)
efC = 200         # construction-time search depth

# IndexHNSWFlat: exact distance computation within HNSW graph
index = faiss.IndexHNSWFlat(DIM, M, faiss.METRIC_INNER_PRODUCT)
index.hnsw.efConstruction = efC

# Embeddings are already L2-normalised, so inner product = cosine sim
index.add(movie_embeddings.astype("float32"))

print(f"Index built. Total vectors: {index.ntotal}")

# Optional: save index to disk
faiss.write_index(index, "movie_hnsw.index")

Step 7 — Real-Time Retrieval

def recommend(user_id: str, top_k: int = 10) -> list:
    """
    Given a user_id, return top-K movie recommendations
    using the Two-Tower model + FAISS HNSW index.
    
    Latency: Query Tower (~2ms) + ANN search (~0.5ms) = ~2.5ms total
    """
    # 1. Run Query Tower: user context → 128-d vector
    user_vector = query_tower.predict(
        {"user_id": tf.constant([user_id])}
    )  # shape: (1, 128)

    # 2. Query FAISS index: find top-K nearest movie vectors
    scores, indices = index.search(
        user_vector.astype("float32"), top_k
    )  # scores: cosine similarity, indices: position in movie array

    # 3. Map indices → movie metadata
    results = []
    for i, (score, idx) in enumerate(zip(scores[0], indices[0])):
        movie_title = movie_titles_all[idx].decode("utf-8")
        results.append({
            "rank":       i + 1,
            "title":      movie_title,
            "score":      round(float(score), 4),
        })
    return results

# Test with User 42
recs = recommend("42", top_k=10)
for r in recs:
    print(f"#{r['rank']:2d} [{r['score']:.4f}]  {r['title']}")

Production Architecture — YouTube vs Spotify vs TikTok

Different products have slightly different constraints that shape how they implement Two-Tower retrieval. Here's how the world's largest recommenders differ:

The Full Production Pipeline

Evaluation — How to Measure Retrieval Quality

Retrieval metrics measure whether the relevant items appear in the top-K candidates. They differ from ranking metrics (which care about the exact order).

Advanced Techniques — Going Beyond the Basics

Golden Rules — Two-Tower in Production