OpenAI Interview: Train a Text-Embedding Model for Search

Part of the OpenAI system design question bank. The question is representative of the round. The analysis and solution are this site's own.

Problem statement

This round is oral. There is usually no diagram tool, and the interviewer often works on search. Expect them to follow you from the training objective into the details: the exact loss, what changes when you double the batch, how negatives are chosen, and then into serving: approximate search, hybrid retrieval, and reranking. Deep RAG product experience is not required; fluency in retrieval fundamentals is.

Clarifying questions

• What is being searched? For practice: a corpus of 100 million passages from documents and help articles, in English.
• What do queries look like? Short, often under ten words, with a mix of natural questions and exact terms.
• What latency? Retrieval under 100 ms at p95, including reranking.
• What labels exist? Search logs with clicks, some question-answer pairs, and a small human-labeled set.
• Is this for RAG or for people? Both. RAG cares most about recall in the top 10 to 20; people care about the top 3.

What makes embedding search hard

Two things pull against each other.

The model must be fast at search time. That forces a bi-encoder: queries and documents embedded separately, documents offline, so search is a nearest-neighbor lookup. But separate encoding means the model never sees the query and the document together, which limits how precisely it can judge relevance.

And the training signal is thin. Most documents are never labeled against most queries. The model learns mostly from what it is told is wrong, the negatives, and the choice of negatives decides what it learns.

So the driving tension is speed versus precision, handled by a pipeline: a fast bi-encoder for recall, and a slow, precise reranker for the final order.

flowchart LR
  Q([Query]):::user --> E[Query encoder]:::svc --> ANN[(ANN index<br/>100 M vectors)]:::store
  ANN -->|top 100| R[Cross-encoder reranker]:::svc -->|top 10| OUT([Results]):::user
  D[Documents]:::store -->|offline| DE[Document encoder]:::svc --> ANN
  classDef user fill:#e6efec,stroke:#315e55,color:#171717;
  classDef svc fill:#f4f1e8,stroke:#315e55,color:#171717;
  classDef store fill:#fdf3dc,stroke:#c4492d,color:#171717;

Key idea. Train the bi-encoder for recall in the top 100 and let a cross-encoder fix the order. Negatives decide what the bi-encoder learns.

Key concepts

Bi-encoder and cross-encoder

A bi-encoder maps a query to a vector and a document to a vector with the same or twin networks. Relevance is their cosine similarity. Documents are embedded once, offline.

A cross-encoder takes the query and document as one input and outputs a score. Every query token can attend to every document token, so it is much more accurate. It must run once per query-document pair at search time, so it only scales to a short candidate list.

Late interaction

A middle ground keeps one vector per token instead of one per document and scores with a sum of maximum similarities between query and document tokens. It is more precise than a single vector and cheaper than a cross-encoder, at the price of a much larger index.

Contrastive learning

The model learns by contrast: pull a query toward its positive document and push it away from negatives. The loss only cares about relative similarity, which is exactly what ranking needs.

Recall at k

The share of queries for which a relevant document appears in the top k results. For a first-stage retriever feeding a reranker, recall at 100 is the metric that matters. The reranker cannot promote a document that was never retrieved.

Key idea. Bi-encoders scale, cross-encoders judge, and contrastive training teaches the bi-encoder what to rank above what.

1. Requirements

Before reading on. What must the retrieval system do, and which metric would you optimize for the first stage?

1.1 Functional requirements

• Train an embedding model from logs and labeled pairs.
• Embed 100 million passages and keep the index fresh as passages change.
• Retrieve candidates for a query with dense and keyword search.
• Rerank candidates and return the top 10.
• Evaluate offline and online, and compare model versions.

1.2 Non-functional requirements

• Recall at 100 above a target, such as 95% on the labeled set.
• Latency. Under 100 ms at p95 end to end.
• Index cost. The vector index must fit in a reasonable memory budget.
• Refresh. New passages searchable within minutes.

1.3 The constraint versus the property

First-stage recall is the property. Anything the retriever misses is lost for good. Latency is the constraint. It forces the bi-encoder, the approximate index, and the short rerank list.

1. Back-of-the-envelope estimation

2.1 Index memory

100 million passages × 768 dimensions × 4 bytes = about 307 GB in float32. In float16, about 154 GB. With product quantization to 64 bytes per vector, about 6.4 GB for codes, plus the graph or list structure. Say that quantization costs some recall and that a rerank stage recovers most of it.

2.2 Embedding the corpus

If one GPU embeds 3,000 passages per second, 100 million passages take about 33,000 seconds, a little over 9 hours on one GPU, or under an hour on 10. Re-embedding the whole corpus for every model change is expensive but feasible. Plan for it, because a new model changes every vector.

2.3 Reranking cost

Reranking 100 candidates per query with a cross-encoder of about 100 million parameters, at about 256 tokens per pair, is 25,600 tokens per query. On a GPU that fits in a few tens of milliseconds when batched. Rerank 1,000 candidates and latency grows tenfold. That is why the list is 100.

2.4 Training examples

Search logs with 50 million clicked query-passage pairs are plenty for fine-tuning. With a batch of 4,096 pairs, one pass over the data is about 12,000 steps.

Key idea. Quantization makes the index affordable, re-embedding is a planned cost, and the rerank list length is set by latency.

1. API design

3.1 Search service

POST /v1/search
  {query, k: 10, filters: {...}, mode: hybrid|dense|keyword}
  -> {results: [{passage_id, score, snippet}], debug?: {dense_rank, keyword_rank, rerank_score}}

3.2 Index maintenance

POST /v1/passages:upsert  [{passage_id, text, metadata}]
POST /v1/passages:delete  [passage_id]
POST /v1/index/versions   {model_version}   -- build a new index for a new model
POST /v1/index/versions/:v:activate

A new embedding model means a new index built side by side, evaluated, then activated. Vectors from two models cannot be mixed.

1. Data model

training_pair   (query, positive_passage_id, source: click|qa|nli|synthetic, weight)
hard_negative   (query, passage_id, miner: bm25|model_v3, teacher_score)
passage         (passage_id, text, metadata, updated_at)
vector_index    (model_version) passage_id -> vector
keyword_index   passage_id -> terms
eval_query      (query, relevant_passage_ids[], graded_relevance[])

1. High-level design

5.1 Keyword search alone

BM25 scores documents by term frequency and rarity. It is fast, needs no training, and wins on exact terms. It fails on meaning: "can I bring my dog to work" does not find the "pet policy" page.

5.2 Fix 1: a bi-encoder trained with in-batch negatives

Train a bi-encoder on query-passage pairs. For each query in a batch, its own passage is the positive and every other passage in the batch is a negative. Search becomes an ANN lookup.

flowchart LR
  B[Batch of B pairs]:::svc --> QE[Encode queries]:::svc
  B --> PE[Encode passages]:::svc
  QE --> S[B x B similarity matrix]:::new
  PE --> S
  S --> L[Softmax over each row:<br/>diagonal is the positive]:::new
  classDef svc fill:#f4f1e8,stroke:#315e55,color:#171717;
  classDef new fill:#ffffff,stroke:#c4492d,stroke-width:2px,stroke-dasharray:5 3,color:#171717;

Meaning-based recall arrives. The model still confuses passages that share a topic, because random in-batch negatives are easy.

5.3 Fix 2: hard negatives

Add, for each query, a few passages that look relevant and are not: top BM25 results or top results from an earlier model that are not the labeled positive. Now the model must learn fine distinctions.

5.4 Fix 3: hybrid retrieval

Run BM25 next to dense retrieval and fuse the lists. Exact-term queries recover their precision.

5.5 Fix 4: a reranker

Rerank the fused top 100 with a cross-encoder. The top 10 improve sharply, which is what users and RAG prompts see.

5.6 The composed system

sequenceDiagram
  autonumber
  participant C as Client
  participant S as Search service
  participant QE as Query encoder
  participant ANN as Vector index
  participant KW as Keyword index
  participant RR as Cross-encoder
  C->>S: query
  par
    S->>QE: embed query
    QE-->>S: vector
    S->>ANN: top 100 by cosine
  and
    S->>KW: top 100 by BM25
  end
  S->>S: reciprocal rank fusion
  S->>RR: rerank fused top 100
  RR-->>S: scores
  S-->>C: top 10

1. Deep dives

6.1 The InfoNCE loss

Before reading on. Write the loss on the whiteboard and say what each term does.

For query q, positive passage p+, and negatives p1 ... pN:

            exp( sim(q, p+) / t )
L = -log  ---------------------------------------------
          exp( sim(q, p+) / t ) + sum_i exp( sim(q, pi) / t )

sim = cosine similarity of L2-normalized vectors
t   = temperature

Read it as a softmax classifier over the candidates: which of these passages is the positive? The loss is the negative log probability of choosing the right one. The gradient pulls q toward p+ and pushes it away from each negative, weighted by how similar that negative currently is. Negatives the model already ranks far below the positive contribute almost nothing. Hard negatives, the ones it confuses, dominate the gradient.

Often the loss is made symmetric: also classify, for each passage, which query it belongs to, and average the two directions.

6.2 What changes when you double the batch

Before reading on. You double the batch size from 4,096 to 8,192. What happens to the loss, the gradient, and the final model?

With in-batch negatives, each query now faces 8,191 negatives instead of 4,095. The task is harder, so the loss value rises at the same model quality. The chance that some negatives in the batch are hard rises, so the gradient carries more useful signal. Empirically, contrastive models improve with larger batches, which is why training spreads a batch across many GPUs and gathers embeddings across devices before computing the similarity matrix.

The costs: memory for the B × B similarity matrix, communication to gather embeddings, and a higher chance of false negatives, pairs in the batch that are actually relevant to each other.

6.3 Hard negatives and false negatives

Before reading on. You mine the top 20 BM25 results as hard negatives. Some of them are actually relevant. What happens, and how do you fix it?

The model is told a relevant passage is wrong and pushes it away. With many such false negatives, recall drops on exactly the queries where several passages are correct.

Fixes, in order of strength:

• Skip the top few mined results, which are the most likely to be true positives.
• Score mined negatives with a cross-encoder teacher and drop any that score close to the positive.
• Use the teacher's scores as soft labels: distill the cross-encoder into the bi-encoder, so a passage can be partly relevant instead of forced to zero.

Refresh mined negatives as the model improves. Negatives mined by an old model become easy for the new one.

6.4 Temperature

The temperature t scales similarities before the softmax. A low temperature, around 0.01 to 0.05, sharpens the distribution: the loss concentrates on the few hardest negatives, which helps separate near-duplicates but can destabilize training and amplify false negatives. A high temperature, near 1, spreads attention over all negatives and learns coarse structure slowly. Many recipes learn t as a parameter or tune it on the validation set.

6.5 Where positive pairs come from

A common recipe pre-trains on large, noisy pairs, then fine-tunes on cleaner, in-domain pairs with hard negatives.

6.6 Serving: ANN indexes

Before reading on. How does HNSW find neighbors without comparing against every vector, and what are its knobs?

HNSW builds a layered graph. Each vector links to a small number of near neighbors. Upper layers are sparse, like an express route; lower layers are dense. A search starts at the top, walks greedily toward the query, drops a layer, and repeats. It visits a tiny fraction of vectors.

Two knobs matter. The number of links per node raises recall and memory. The search breadth at query time raises recall and latency. Tune them to hit the recall target at the latency budget.

For very large corpora, inverted-file indexes with product quantization cut memory further: cluster the vectors, search only the nearest clusters, and store compressed codes. Recall drops a little, and reranking with the full vectors or the cross-encoder recovers it.

6.7 Hybrid fusion

Dense and BM25 scores live on different scales, so adding them needs tuning. Reciprocal rank fusion avoids that: score each passage by the sum of 1 / (k + rank) over the lists it appears in, with k around 60. A passage ranked high in either list rises; one high in both rises most. It needs no calibration, which is why it is the common default.

6.8 Reranking: pointwise, pairwise, listwise

Before reading on. Compare the three families. When would you move from pointwise to listwise?

Have this ready before the question lands. A vague answer here is a reported place to lose the round.

A practical answer: start with a pointwise cross-encoder, because it is strong and cheap to train. Move to listwise when the top few positions carry most of the value, such as a RAG prompt that takes only five passages, and when features beyond text matter: freshness, popularity, authority. Gradient-boosted rankers with listwise objectives remain common there. An LLM can rerank listwise over 20 passages at a time, with a sliding window for longer lists, at a much higher cost.

6.9 Evaluation

Offline, on a held-out labeled set: recall at 100 for the first stage, nDCG at 10 and MRR for the final list. Slice by query type: exact-term queries, long natural questions, rare topics. Averages hide the slices that users complain about.

Online: click-through, time to a successful click, and the reformulation rate, which rises when results are bad. For RAG, measure answer quality downstream, because better retrieval only matters if answers improve. Run A/B tests before switching models.

6.10 The training loop, in code

Before reading on. Write the core of one training step for a bi-encoder with in-batch negatives.

import torch
import torch.nn.functional as F

def training_step(query_encoder, passage_encoder, queries, positives, hard_negs, tau=0.05):
    # queries: B inputs; positives: B inputs; hard_negs: B*H inputs
    q = F.normalize(query_encoder(queries), dim=-1)          # B x d
    p = F.normalize(passage_encoder(positives), dim=-1)      # B x d
    n = F.normalize(passage_encoder(hard_negs), dim=-1)      # (B*H) x d

    candidates = torch.cat([p, n], dim=0)                    # (B + B*H) x d
    logits = q @ candidates.T / tau                          # B x (B + B*H)
    labels = torch.arange(q.size(0), device=q.device)        # positive i sits at column i
    return F.cross_entropy(logits, labels)

Walk through it. Each query's positive is at its own index in the candidate list; every other positive in the batch, and every mined hard negative, serves as a negative. Cross-entropy over each row is exactly the InfoNCE loss. Doubling B adds columns, which means more negatives per query, which is the effect discussed in section 6.2.

In multi-GPU training, gather p and n from all devices before building candidates, so each query sees negatives from the whole global batch. That is where the memory and communication costs come from.

6.11 Choosing the embedding dimension

A larger dimension stores more information per vector and usually improves recall slightly. It also grows the index linearly and slows every distance computation. For practice: going from 768 to 1,536 dimensions doubles the float16 index for 100 million passages from about 154 GB to about 307 GB.

Two techniques relax the tradeoff. Train with nested objectives so the first 256 dimensions are useful on their own, then store short vectors for the first stage and use the full vector to re-score candidates. And quantize: 8-bit or even binary codes for the first stage, with full-precision re-scoring of the top few hundred. Measure recall at 100 for each option on the same evaluation set; pick the smallest representation that meets the recall target.

1. Variants

7.1 Multilingual search

Train on parallel data so a query in one language finds passages in another. Check recall per language; high-resource languages mask weak ones in the average.

7.2 Code search

Code has exact identifiers and structure. Keyword search carries more weight, and chunking by function beats fixed windows.

7.3 Domain adaptation without labels

Generate synthetic queries for in-domain passages with an LLM, filter them with a cross-encoder, and fine-tune on them. It often beats a general model on specialized corpora.

7.4 At ten times the corpus

At a billion passages, the float16 index alone is about 1.5 TB. Quantize the first stage aggressively and shard the ANN index across machines by passage ID, querying all shards in parallel and merging. Re-embedding the corpus for a new model becomes a multi-day batch job, so run it in the background, build the new index side by side, and switch only after evaluation. Keep the old index until the new one has served real traffic without regression.

1. The transferable pattern

Embedding search is cheap recall, then expensive precision. A fast model narrows millions to a hundred, and a slow model orders the hundred. The same cascade runs recommendation systems, ad ranking, and fraud detection. The transferable skill is knowing which stage owns which metric: recall for the first, precision at the top for the last.

Review: the 30-second answer

• Bi-encoder for recall, cross-encoder for order. Optimize recall at 100 first.
• InfoNCE with in-batch negatives. Bigger batches give more and harder negatives.
• Mine hard negatives, filter false ones. Use a cross-encoder teacher, and distill it.
• HNSW or IVF-PQ for serving. Tune links and search breadth to the latency budget.
• Hybrid with rank fusion, then rerank. Pointwise by default, listwise when the top few positions are what matter.

Quiz

Sources and further reading

• Representation Learning with Contrastive Predictive Coding introduced the InfoNCE loss.
• Dense Passage Retrieval for Open-Domain Question Answering trains a bi-encoder with in-batch and BM25 hard negatives.
• SimCSE: Simple Contrastive Learning of Sentence Embeddings covers dropout-based positives and NLI-based hard negatives.
• ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction describes late interaction.
• HNSW describes the graph index.
• The full product built on this retrieval stack is the enterprise RAG assistant.