Two-Tower Search

Week 2 of the residency. TF-IDF, PageRank, and then a search system built out of the word2vec model from week 1.

Week 2: search. Given a query, find the document that answers it.

Strip away the machinery and every search system is one function. A document goes in, a query goes in, a number comes out.

Everything in this week is a different way of filling that box in. Rank five documents by running the function five times and sorting.

TF-IDF

TF-IDF fills the box with counting. Term frequency is how often a query word appears in a document. Inverse document frequency discounts words that appear everywhere: if a term shows up in every document it cannot distinguish between them, so its weight goes to zero. Multiply the two and you have a score, with no learning anywhere.

It also builds an index the other way round, from term to the documents containing it, so answering a query is a lookup rather than a scan of the corpus. That is how search engines worked for decades, and partly still do.

Is It Good Enough?

No, and the reason is in the paper that founded Google:

On the web, this strategy often returns very short documents that are the query plus a few words. For example, we have seen a major search engine return a page containing only “Bill Clinton Sucks” and picture from a “Bill Clinton” query. Some argue that on the web, users should specify more accurately what they want and add more words to their query. We disagree vehemently with this position.

— Brin and Page, The Anatomy of a Large-Scale Hypertextual Web Search Engine, 1998

Counting terms rewards a page for containing the query words. It cannot tell a good page from a bad one that uses the same vocabulary.

The fix is to score the document on its own, separately from the query, using how the rest of the web treats it.

PageRank

Documents linking to each other form a graph, and a graph is an adjacency matrix. Normalise the rows and that matrix becomes a set of transition probabilities: if you are on page A, where do you go next?

That is a random walk. The teleportation term, α = 0.15, is the chance the walker ignores the links and jumps to a random page instead. Without it the walk can get stuck in a corner of the graph with no way out.

Run the walk forever and the fraction of time spent on each page settles down. That settled distribution is the answer to π = πT, and it is PageRank: a page is important if a random surfer spends a lot of time there.

The part I found genuinely clarifying is that none of this was new mathematics. Markov described these walks in 1906. PageRank is a standard stationary distribution applied to the link graph, and the teleportation term is there to guarantee the walk has one.

The Other Half

PageRank scores a document on its own merit. It says nothing about whether the document answers this query, and that is still TF-IDF’s job, still by counting words.

So: can we do better than TF-IDF?

The problem with counting is that a query for was ronald reagan a democrat will not match a passage that answers the question without using those exact words. Matching on literal overlap misses meaning.

Week 1 had the same problem, and week 1 already built the answer to it.

CBOW Is an Encoder

This is the slide the whole week turns on.

The CBOW model from week 1 was built to predict a missing word. But look at what it does on the way: it takes a sentence, looks up a vector per word, and averages them into one vector. That intermediate step was a means to an end, and it is also a complete answer to “turn this sentence into numbers”.

Word2vec is an encoder. It was one all along.

So the relevance half of search is now: encode the query, encode the document, and compare the two vectors. Similar meaning lands in a similar place, whether or not the same words were used.

Two Towers

One encoder for the query, one for the document, meeting only at the score. They never see each other.

That constraint is what makes it usable. The document tower never sees a query, so it can run before any query exists.

• Embed the whole corpus once, offline
• Put the vectors in an index
• At query time, embed one short string and look up its nearest neighbours

The alternative is one model reading the query and document together. It is more accurate. It also has to run once per document per query, which does not survive contact with a real corpus.

And because the towers only have to agree on a shared vector space, they do not have to be the same kind of model.

A CNN trained to classify images produces exactly the same kind of object partway through: a vector that summarises its input.

Swap the document tower for the image one and the architecture is unchanged. You can now search photographs with a sentence.

Nothing about the design had to be rethought. This is the first time in the course that a piece of machinery gets reused for something it was not built for. It is not the last.

Building One

The rest of this page trains a small two-tower model, because the part that decides whether it works is not the architecture.

There is a trick for getting query-document pairs without labelling anything: take a document, cut a piece out, and use the piece as the query. Here I take Hacker News titles of at least eight words and cut them in half.

import gzip
import json
import re
import sys
import zlib

sys.path.insert(0, "..")
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from _style import COLOURS, MUTED, figure, style_axes

REVISION = "f705fed08827ff6c36e3b5329495c943a5e544e8"
path = hf_hub_download("roshbeed/ai-residency-blog-data", "hn/hn-sample.json.gz",
                       repo_type="dataset", revision=REVISION)
rows = json.load(gzip.open(path, "rt"))
titles = [r["title"] for r in rows["train"]] + [r["title"] for r in rows["test"]]

tokenise = lambda s: re.findall(r"[a-z0-9\+#]+", s.lower())
pairs = []
for title in titles:
    words = tokenise(title)
    if len(words) >= 8:
        half = len(words) // 2
        pairs.append((words[:half], words[half:]))

print(f"{len(pairs):,} query/document pairs")
print(f"  query:    {' '.join(pairs[0][0])}")
print(f"  document: {' '.join(pairs[0][1])}")

79,009 query/document pairs
  query:    california governor to deploy 500 surveillance
  document: cameras to oakland to fight crime

Each half of a title becomes a bag of hashed words, and the two towers turn those into unit vectors. Recall@10 asks how often the right document lands in the top ten of all five thousand, and chance is the number to hold it against.

DIM = 2**14


def encode(sequences):
    """Bag of hashed words. crc32 rather than hash() — Python's is salted per process."""
    X = torch.zeros(len(sequences), DIM)
    for i, words in enumerate(sequences):
        for w in words:
            X[i, zlib.crc32(w.encode()) % DIM] += 1.0
        if words:
            X[i] /= len(words)
    return X


rng = np.random.default_rng(0)
order = rng.permutation(len(pairs))
train_ids, test_ids = order[:40_000], order[40_000:45_000]

Q, D = encode([pairs[i][0] for i in train_ids]), encode([pairs[i][1] for i in train_ids])
Q_test = encode([pairs[i][0] for i in test_ids])
D_test = encode([pairs[i][1] for i in test_ids])


class Tower(nn.Module):
    """Text in, a unit vector out. The query and document sides get one each."""

    def __init__(self, out=128):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(DIM, 256), nn.ReLU(), nn.Linear(256, out))

    def forward(self, x):
        return F.normalize(self.net(x), dim=-1)


def recall_at_10(query_tower, doc_tower):
    with torch.no_grad():
        similarity = query_tower(Q_test) @ doc_tower(D_test).T
        top = similarity.topk(10, dim=1).indices
        return (top == torch.arange(len(Q_test)).unsqueeze(1)).any(1).float().mean().item()


CHANCE = 10 / len(D_test)
print(f"{len(Q):,} training pairs, {len(D_test):,} documents to search")
print(f"chance recall@10 = {CHANCE:.4f}")

40,000 training pairs, 5,000 documents to search
chance recall@10 = 0.0020

Drawn out, one tower is a stack of ordinary layers:

from _arch import diagram

diagram(Tower(), input_shape=(1, DIM))

Figure 1: One tower. The query and document sides are each one of these, with separate weights, and the output is normalised so a dot product between two of them is a cosine.

Both towers start random, so their shared space means nothing yet. The training signal has to teach them what close means, and the obvious way is a triplet: a query, a document that answers it, and one that doesn’t. Pull the first pair together, push the second apart, by at least a margin.

def train(loss_fn, epochs=8, batch=256, seed=0):
    torch.manual_seed(seed)
    query_tower, doc_tower = Tower(), Tower()
    optimiser = torch.optim.Adam([*query_tower.parameters(), *doc_tower.parameters()], lr=1e-3)
    generator = torch.Generator().manual_seed(seed)

    losses, recalls = [], []
    for _ in range(epochs):
        perm = torch.randperm(len(Q), generator=generator)
        total = steps = 0
        for i in range(0, len(perm) - batch, batch):
            b = perm[i:i + batch]
            loss = loss_fn(query_tower, doc_tower, b, perm, generator)
            optimiser.zero_grad()
            loss.backward()
            optimiser.step()
            total += loss.item()
            steps += 1
        losses.append(total / steps)
        recalls.append(recall_at_10(query_tower, doc_tower))
    return losses, recalls


MARGIN = 0.2


def triplet(query_tower, doc_tower, b, perm, generator):
    """One document picked at random as the negative."""
    negatives = perm[torch.randperm(len(b), generator=generator)]
    q = query_tower(Q[b])
    positive = (q * doc_tower(D[b])).sum(-1)
    negative = (q * doc_tower(D[negatives])).sum(-1)
    return F.relu(MARGIN - positive + negative).mean()


triplet_loss, triplet_recall = train(triplet)
for epoch, (l, r) in enumerate(zip(triplet_loss, triplet_recall), 1):
    print(f"epoch {epoch}: loss {l:.4f}   recall@10 {r:.4f}   ({r / CHANCE:.1f}x chance)")

epoch 1: loss 0.0121   recall@10 0.0020   (1.0x chance)
epoch 2: loss 0.0093   recall@10 0.0022   (1.1x chance)
epoch 3: loss 0.0123   recall@10 0.0022   (1.1x chance)
epoch 4: loss 0.0199   recall@10 0.0020   (1.0x chance)
epoch 5: loss 0.0142   recall@10 0.0020   (1.0x chance)
epoch 6: loss 0.0201   recall@10 0.0018   (0.9x chance)
epoch 7: loss 0.0179   recall@10 0.0020   (1.0x chance)
epoch 8: loss 0.0178   recall@10 0.0018   (0.9x chance)

The loss falls to near zero within one epoch and recall never leaves chance. Out of 5,000 documents the model finds the right one in its top ten about as often as picking ten at random.

A negative drawn at random is about a completely different subject, so the towers only have to tell machine learning from sourdough. They manage that immediately, the margin is satisfied, and a satisfied margin has no gradient. The model stops learning while still unable to do the thing you want, which is to pick the right passage out of a thousand plausible ones.

The standard fix is to use the whole batch: every other document in it is a negative for this query, so the question stays hard as the model improves, and the documents are already encoded.

TEMPERATURE = 0.05


def in_batch(query_tower, doc_tower, b, perm, generator):
    """Every other document in the batch is a negative for this query."""
    q, docs = query_tower(Q[b]), doc_tower(D[b])
    return F.cross_entropy((q @ docs.T) / TEMPERATURE, torch.arange(len(b)))


batch_loss, batch_recall = train(in_batch)
for epoch, (l, r) in enumerate(zip(batch_loss, batch_recall), 1):
    print(f"epoch {epoch}: loss {l:.4f}   recall@10 {r:.4f}   ({r / CHANCE:.1f}x chance)")

epoch 1: loss 5.3606   recall@10 0.0102   (5.1x chance)
epoch 2: loss 4.8786   recall@10 0.0158   (7.9x chance)
epoch 3: loss 4.0110   recall@10 0.0258   (12.9x chance)
epoch 4: loss 2.4485   recall@10 0.0324   (16.2x chance)
epoch 5: loss 0.9365   recall@10 0.0352   (17.6x chance)
epoch 6: loss 0.3389   recall@10 0.0366   (18.3x chance)
epoch 7: loss 0.1411   recall@10 0.0370   (18.5x chance)
epoch 8: loss 0.0765   recall@10 0.0364   (18.2x chance)

Both runs plotted together, loss on the left and recall on the right.

import matplotlib.pyplot as plt

epochs = range(1, len(batch_loss) + 1)
fig, (left, right) = plt.subplots(1, 2, figsize=(8.4, 3.8))

left.plot(epochs, triplet_loss, color=COLOURS[1], linewidth=2)
left.plot(epochs, batch_loss, color=COLOURS[0], linewidth=2)
left.set_title("training loss", fontsize=10, color=MUTED, pad=8)
style_axes(left, "Epoch", "Loss")

right.axhline(CHANCE, color=MUTED, linewidth=1.2, linestyle="--")
right.plot(epochs, triplet_recall, color=COLOURS[1], linewidth=2)
right.plot(epochs, batch_recall, color=COLOURS[0], linewidth=2)
right.annotate("in-batch negatives", xy=(epochs[-1], batch_recall[-1]), xytext=(-6, -14),
               textcoords="offset points", ha="right", fontsize=9, color=COLOURS[0])
right.annotate("one random negative", xy=(epochs[-1], triplet_recall[-1]), xytext=(-6, 8),
               textcoords="offset points", ha="right", fontsize=9, color=COLOURS[1])
right.annotate("chance", xy=(1, CHANCE), xytext=(2, 6), textcoords="offset points",
               fontsize=9, color=MUTED)
right.set_title("recall@10 out of 5,000 documents", fontsize=10, color=MUTED, pad=8)
style_axes(right, "Epoch", "Recall@10")
fig.tight_layout()

Figure 2: The same towers, the same data, the same number of steps. Only the choice of negatives differs.

Same towers, same data, same number of gradient steps. The triplet run even ends with the lower training loss. Rank the two by that number and you pick the one that learned nothing, which is why the chance line is on the right-hand panel.

Conclusion

Week 2 builds a search system in two halves.

• PageRank scores a document on its own merit, from the link graph
• A learned encoder scores whether it is about the right thing
• The encoder is week 1’s model, reused
• The same two-tower design works for images with no changes

The second half only exists because week 1 built an encoder while trying to do something else. That reuse is the thread running through the rest of the course.

The full project, with GRU towers and a Redis vector index, is on GitHub.