Embedding-based retrieval in Facebook search

01 · CONTEXT

Why social search calls for embeddings

At the time of the paper, Facebook search still relied mainly on a Boolean matching model, although embedding-based retrieval had already been used in web search engines for years. There, search intent does not depend on the text alone: it is also influenced by who is searching and the context in which the search happens. For that reason, according to the authors, embedding-based retrieval at Facebook is not a text embedding problem.

A search has a stage of retrieval, which looks for a set of relevant documents with low latency and low computational cost, and a ranking stage, which puts the most desirable documents on top with more complex models. The difference in scale is large: ranking usually considers hundreds of documents per session, while retrieval has to process billions or trillions of documents in the index.

Embedding-based retrieval (EBR) uses embeddings to represent the query and the documents and converts the retrieval problem into a nearest neighbor search in the embedding space.

02 · MODEL AND TRAINING

Two-sided unified embedding

The authors formulate retrieval as a recall optimization problem. The target result set for a query can come from clicks or from documents judged relevant by human evaluation. As an offline metric, they propose running a KNN search over the entire index and using recall@K.

O modelo is unified and has two sides: on one, the search request, with query text, the person searching, and context; on the other, the document. The two encoders are, by default, separate networks, with the option of sharing part of the parameters. The dense vectors are compared by cosine similarity, and the intuition of the triplet loss is to separate the positive pair from the negative pair by a distance margin.

The margin requires tuning: the optimal value varies considerably across training tasks, and different margins resulted in a 5-10% variation in KNN recall.

Inputs beyond text make a difference. Swapping text embeddings for unified embeddings improved recall by +18% in event search and +16% in group search. Among the added attributes are location, advantageous in searches for local businesses, groups, and events, and social attributes, coming from a separate embedding model trained on the social graph.

Compared with Boolean term matching, embeddings trained only on text attributes proved good in two scenarios: approximate text matching and term optionalization, such as dropping a term from the query and still retrieving the expected group.

The authors attribute the regression to a bias: the no-click negatives tend to be hard cases, while most documents in the index are easy cases, which do not match the query. Training only with negatives like that changes the representativeness of the training data relative to the real retrieval task.

03 · SERVING

Hybrid retrieval over the inverted index

The direct way to serve would be to merge the candidates from the existing retrieval with those from an embedding KNN search. The authors considered this suboptimal for three reasons: high performance cost in the initial experiment, high maintenance cost from having two indexes, and large possible overlap between the two candidate sets.

The solution was hybrid retrieval, which integrates embedding KNN and Boolean matching to score documents. For this they used the Faiss library for vector quantization and integrated it into inverted index retrieval. The authors cite two advantages of inverted-index-based approximate search: lower storage cost, through vector quantization, and easier integration with the existing retrieval system.

The integration was built in Unicorn, the retrieval engine behind most of Facebook's search products. The authors extended the document representation to include embeddings and created the query operator (nn), which matches the documents whose embedding is within a radius of the query embedding. At indexing time, each document's embedding is quantized and becomes a term, for its coarse cluster, and a payload, for the quantized residual. At query time, the operator is rewritten internally as an OR of the terms of the clusters closest to the query embedding.

The two-sided model was decomposed into a query model and a document model, served separately. The query goes through an online inference service, in real time. Documents are processed offline in batch with Spark, and the generated embeddings are published, along with other metadata, to the forward index; with the additional quantization (coarse and PQ), also to the inverted index.

1. The query arrives with the text, the searcher, and the context, and goes to the online service that computes the request embedding in real time.

2. The query embedding goes on to the hybrid retrieval engine.

3. The engine rewrites the nearest-neighbor operator as an OR of the terms of the closest clusters and combines it with the Boolean terms of the expression.

4. Off the query path, batch inference publishes the already-quantized document embeddings to the index.

5. The retrieved candidates move on to the ranking and filtering stages.

Because nearest-neighbor support was built on existing primitives rather than as a separate system, the engine inherited real-time updates, efficient query planning and execution, and multi-hop queries. Among the search modes, radius search gave a better balance between system performance and quality than top-K, and it is the one used in the production setup described in the paper.

04 · ANN TUNING

Quantization and approximate search parameters

Embedding quantization has two components: coarse quantization, which groups the vectors into clusters, typically with K-means, and product quantization, which is finer and allows distances to be computed efficiently. Among the parameters to tune are the number of coarse clusters, which affects performance and recall, the number of PQ bytes, and nprobe, which defines how many clusters will be scanned and also affects performance and recall.

The authors also recommend deploying several online ANN configurations to understand the real performance impact of embedding-based retrieval on the system, which helps decide the capacity budget and narrow the scope of the offline parameter search.

05 · HARD MINING

Hard negatives and positives

In the people search model, the top K results by embedding generally had the same name, and the model did not always rank the target above the others, even with social attributes. The authors considered it very likely that the training negatives were too easy, since they were random samples, generally with different names. To help the model better distinguish similar results, they proposed using samples closer to the positives in the embedding space as hard negatives.

In online HNM, hard negatives are selected dynamically in each training batch: the documents with the highest similarity become the hardest negatives. Enabling it improved recall by +8.38% for people, +7% for groups, and +5.33% for events. The optimum was at most two hard negatives per positive, and using more than that starts to make the model worse. The limitation is that the chance of having a hard negative among random samples can be low.

1. Search logs feed example mining, which at the starting point uses clicks as positives and random documents as negatives.

2. The two-sided model is trained with the triplet loss, which separates the positive pair from the negative one by a distance margin.

3. In offline HNM, the model generates the top results for each query, and approximate search greatly reduces the compute time of that generation.

4. Hard negatives are sampled from an intermediate band of the ranking, not from the hardest examples.

5. The model is retrained mixing easy and hard negatives, increasing the proportion of easy ones relative to hard ones.

6. After training, the model is split into a query encoder and a document encoder, served separately.

Offline HNM generates the top K results for each query, selects hard negatives by some strategy, and retrains the model with the new triplets. One result that seems counterintuitive: models trained only with hard negatives did not outperform those trained with random negatives. Using the hardest examples was not the best strategy either, and sampling among ranking positions 101-500 gave the best model recall.

Mixing random and hard negatives helped: increasing the proportion of easy ones kept improving recall, saturating at 100:1. Transfer learning from the hard model to the easy one brought an additional gain. To generate the top results without an exhaustive search, approximate search greatly reduces compute time, and running it on a random shard was enough to generate effective negatives.

On the positives side, the authors looked for new results that production did not yet retrieve successfully, mining potential targets from failed search sessions. A model trained only on these hard positives reached recall similar to the click data with only 4% of the data volume.

06 · AFTER RETRIEVAL

Selection, ranking, ensemble, and result

To improve efficiency and quality, the authors applied query and index selection. EBR is not triggered for certain queries where it would be weak and add no extra value. In the index, they chose, for example, only users active in the month, recent events, and popular pages and groups.

Because the existing ranking stages were designed for the previous retrieval scenarios, the authors brought embedding similarity into ranking as a feature, and cosine similarity performed best among the options tested. Because EBR can have lower precision than term matching, they also built a feedback loop with human evaluation to retrain the relevance model and filter out irrelevant results.

The HNM analysis led to the ensemble idea: hard examples are needed to improve precision, but easy ones also matter for representing the retrieval space. Hence the multi-stage approach, in which the first model focuses on recall and the second specializes in telling apart more similar results. Among the options tested for the second stage, the model trained with non-click negatives had the best KNN recall, but suffered much more precision loss when quantized.

Another cascade composition used the text embedding to pre-select candidates that match the text and then the unified model to rerank and return the best candidates. The authors report significant metric gains in online A/B tests in the verticals where they evaluated EBR.