Embeddings and similarity search

A well-run clothes shop arranges its floor so that similar things sit near each other. Plain white shirts are next to striped shirts, which are next to light jumpers, and a customer who holds up a shirt and asks for ‘something like this’ is walked two metres along the rail. The person who decided where everything goes had to settle what ‘like this’ means, by colour, by occasion or by price, and every later recommendation depends on that choice.

An embedding is a position on that floor, except the floor has a few hundred dimensions instead of two and the positions are learned from data rather than decided by a manager. The model that assigns positions is the encoder, the floor is the embedding space, and ‘walk two metres along the rail’ is a nearest neighbour search, meaning the lookup of the items whose positions are closest to the one in hand.

Size is where the shop floor stops helping, because a shop floor has two dimensions and a few thousand items, so you can see the neighbours at a glance. A catalogue has 10 million items with 128 numbers each, and checking the distance to every one of them for every query is the cost this topic has to design around. The running example is visual search in a shopping app, where a customer photographs a jacket and wants to buy one like it.

The encoder is the model that turns an image, a product or a user into its list of numbers, and what it was trained on decides what ‘close’ means. In visual search, finding the exact same jacket from a different photo and finding a jacket in the same style are different tasks, and an encoder trained for one will disappoint on the other. The training data is pairs, two photos of the same jacket for the first task or two jackets a stylist calls similar for the second, together with unrelated jackets that the encoder is taught to push apart.

A two-tower model is the common shape for matching a query to items. One encoder handles the query, which for visual search is the customer’s photo, and a second encoder handles the catalogue items, and training pushes the two lists together for matching pairs and apart for the unrelated ones. The split matters for serving because the item side does not depend on the request, so all 10 million item lists can be computed overnight and stored, and at request time only the photo is encoded. Ten million lists of 128 numbers at four bytes each is about 5.1 gigabytes, which fits in the memory of one serving machine.

A similarity function is the rule for turning two lists into one closeness number. The three usual ones are called dot product, cosine similarity and Euclidean distance, and each is a different recipe for that number. They do not rank the same neighbours, because the length of a list affects two of them and not the third. If the encoder was trained with lists scaled to length one, then the dot product and the Euclidean distance agree, and that scaling has to happen at search time too. An index that measures distance differently from the training rule returns neighbours the encoder never meant, like a shop assistant measuring the rail in a different unit from the person who laid it out.

Exact search compares the query with every stored list, and with 10 million items of 128 numbers that is about 1.3 billion multiplications per query, which is a few hundred milliseconds on one processor and too slow at checkout. Approximate nearest neighbour search, usually shortened to ANN, organises the lists into regions and only searches the few regions nearest the query, and it finishes in a few milliseconds at the price of sometimes missing a true neighbour. The quality of the shortcut is measured as ANN recall, the share of the true top ten that the approximate search also returned, and 95 percent is a common target.

Filters remove results after the search has done its work. If the search returns the 100 closest jackets and 95 of them are out of stock in the customer’s size, the customer sees five results, however good the search was. The remedies are to filter inside the index so it skips ineligible items as it searches, to retrieve more candidates than you need and filter afterwards, or to keep a separate index per large category. Each needs testing on the rare sizes and the newest items, where the filter removes the most.

A new version of the encoder rearranges the floor. It may keep the same 128 numbers per item while putting shirts where jumpers used to be. A query encoded with the new encoder and searched against an index built by the old one then returns neighbours that are nonsense, with no error anywhere. The encoder and the index it built are one versioned pair, and a deployment replaces both at once or runs the new pair alongside the old one until it is checked.

The index also has to cope with change between rebuilds. New items arrive hourly, items sell out, and some are withdrawn for legal reasons that take effect immediately. The design therefore needs an incremental path that adds and removes a few items without a full rebuild, plus an eligibility check at serving time for anything that must not appear regardless of what the index says. With the floor laid out and searchable, the next topic is what to do for the customer or item that has no history to be placed by, which is the cold start problem.

A high ANN recall gets read as proof that results are good. ANN recall says the approximate search found the same neighbours as the exact search in the same space, and if the encoder learned the wrong idea of similar, both searches agree on the wrong neighbours.

The similarity function also gets mixed up between training and serving. An encoder trained with cosine similarity and an index using the raw dot product will rank long lists above close ones, and the symptom is a handful of items that appear as neighbours to everything.

Updating the encoder like any other model is the confusion that costs a full rebuild. Every other model can be swapped behind the same inputs, while an encoder changes the meaning of the stored data it produced, so the stored data has to be regenerated with it.

In any search or recommendation design, the modelling block spends three to five minutes on embeddings and retrieval. Introduce them with the two-tower split and the reason for it, for example ‘I would encode items overnight into an index and encode the query at request time. Then I would search with an approximate nearest neighbour index that trades a little recall for a lookup in a few milliseconds’. The follow-ups are almost always about filters, the retrain, or what similar means, and each has a section above.