Dense Retrieval: Semantic Search & Bi-Encoder Implementation › The Bi-Encoder Architecture

The dominant architecture for dense retrieval is the bi-encoder, which uses two separate encoder networks to produce embeddings for queries and documents independently. This separation is important for efficiency at scale, and understanding why requires us to think carefully about what happens during retrieval.

Think of a bi-encoder as two photographers: the query encoder takes a photograph of the question, and the document encoder takes a photograph of each document. The photographs are all printed in the same format, so you can compare them side by side using a simple similarity measure. Because the document photographs can be taken before any query arrives (offline), you only need to take the query photograph at search time and then compare it against the pre-existing document photographs. This is what makes the bi-encoder scalable: the expensive work of encoding documents happens once, not once per query.

The bi-encoder consists of two components that work in tandem to transform text into comparable numerical representations. The first component is the query encoder, denoted $E_q$, which takes a natural language query $q$ as input and produces a dense vector representation. The second component is the document encoder, denoted $E_d$, which performs the analogous transformation for documents.

Formally, these transformations are:

• Query encoder $E_q$: Maps a query $q$ to a dense vector $\mathbf{q} = E_q(q) \in \mathbb{R}^D$
• Document encoder $E_d$: Maps a document $d$ to a dense vector $\mathbf{d} = E_d(d) \in \mathbb{R}^D$

Both encoders are typically initialized from the same pre-trained transformer (such as BERT, as we covered in Part XXIV), and they may share weights or be fine-tuned independently. The encoders produce embeddings of the same dimensionality $D$, letting direct comparison. This shared dimensionality is needed because it allows us to measure distances and angles between query and document vectors in the same geometric space. Without a shared dimensionality, there would be no principled way to compare a query vector to a document vector.

In practice, the query and document encoders are often initialized identically and may even share weights during training. Weight sharing reduces the total parameter count and can improve training stability. Some architectures allow the query and document encoders to diverge during fine-tuning, which can be beneficial when the distributional characteristics of queries and documents are very different. Short, keyword-like queries and long, paragraph-length documents are stylistically distinct, and separate encoders give the model flexibility to specialize. For many practical applications, however, a single shared encoder works well and is simpler to maintain.

A key challenge in building these encoders is converting the variable-length output of a transformer into a single fixed-size vector suitable for comparison. Recall from our BERT chapters that when a BERT model processes an input sequence, it produces a hidden state vector for each token in the sequence. A sentence of 20 words produces 20 hidden state vectors (plus special tokens), while a paragraph of 200 words produces 200 hidden state vectors. We need to compress this variable-length sequence into a single fixed-size vector, a process called pooling.

BERT-based encoders typically extract the representation of the special [CLS] token from the final layer:

$$\mathbf{q} = \text{BERT}(q)_{[\text{CLS}]}$$

where:

• $\mathbf{q}$ is the dense vector representation of the query
• $\text{BERT}(q)$ is the sequence of hidden states from the last layer of the BERT model
• $[\text{CLS}]$ indicates extraction of the vector corresponding to the special classification token, which is designed to aggregate sequence-level information during pre-training

The [CLS] token was specifically designed during BERT's pre-training to aggregate information from the entire sequence. During the next-sentence prediction pre-training task, the [CLS] representation is used to classify whether two sentences are consecutive, forcing it to encode sequence-level rather than token-level information. This makes it a natural candidate for creating a single vector that summarizes the entire input.

Alternatively, to capture information distributed across the entire sequence rather than relying on a single token, some models compute the average of all token representations. This approach, known as mean pooling, treats each token's contribution equally:

$$\mathbf{q} = \frac{1}{n} \sum_{i=1}^{n} \text{BERT}(q)_i$$

where:

• $\mathbf{q}$ is the dense vector representation of the query
• $n$ is the number of tokens in the query
• $\text{BERT}(q)_i$ is the vector representation of the $i$-th token
• $\sum$ computes the element-wise sum of all token vectors, and dividing by $n$ finds the geometric centroid

The intuition behind mean pooling is that important semantic information may be spread across multiple tokens rather than concentrated in the [CLS] token. By averaging, we create a representation that balances contributions from all parts of the input. Consider a sentence like "The quick brown fox jumped over the lazy dog." The [CLS] representation tries to summarize this whole sentence, but mean pooling creates a vector that represents the average meaning of all the individual tokens. For retrieval purposes, mean pooling often works slightly better because it ensures no important content is lost.

In practice, the choice between [CLS] pooling and mean pooling affects retrieval quality, and different models adopt different strategies based on their training objectives. Empirically, models trained with mean pooling as part of their objective tend to perform better when evaluated using mean pooling, and likewise for [CLS] pooling. The Sentence-BERT family of models (which we use in our implementation) were specifically trained with mean pooling and perform substantially better with it.

The bi-encoder's separation of query and document encoding enables a necessary optimization: pre-computation of document embeddings. To appreciate why this matters, consider the computational demands of a retrieval system at real scale.

In a production retrieval system serving a corpus of 10 million documents, we need to compare every query against every potential document to find the most relevant ones. If we process each query-document pair jointly, a single query requires 10 million transformer forward passes. A transformer forward pass on a modern BERT-sized model takes approximately 10 milliseconds on a GPU. This means a single query would take 100,000 seconds, or about 27 hours. That is obviously not a retrieval system; it is a computational nightmare.

With the bi-encoder, we pre-compute document embeddings once. At query time, we run the transformer once for the query (taking 10 milliseconds), then compute dot products between the query vector and all 10 million pre-computed document vectors. Computing 10 million dot products in a 768-dimensional space takes microseconds using vectorized operations. The total query latency drops from 27 hours to roughly 50 milliseconds. This is the bi-encoder's core value proposition: it separates the expensive transformer computation (done once per document, offline) from the cheap similarity computation (done at query time, over pre-computed vectors).

This contrasts with cross-encoders, which concatenate the query and document and process them jointly through a single transformer. Cross-encoders can capture fine-grained interactions between query and document tokens. The attention mechanism allows every query token to attend to every document token, letting the model to reason about how specific parts of the document relate to specific parts of the query. This richer interaction produces more accurate relevance judgments in controlled experiments.

However, cross-encoders require running the transformer once for every query-document pair. For a corpus of 10 million documents, this means 10 million transformer forward passes per query, which is computationally infeasible for first-stage retrieval. Cross-encoders therefore serve a different role in the retrieval pipeline: they are used as rerankers that refine the top candidates returned by a faster first-stage retriever. We will explore reranking with cross-encoders in a later chapter.

The bi-encoder architecture trades some modeling power for massive efficiency gains:

We will explore reranking with cross-encoders in a later chapter, where they serve as a second-stage refinement over bi-encoder candidates. The combination of bi-encoder retrieval followed by cross-encoder reranking is a powerful and widely used two-stage pipeline.

Beyond [CLS] and mean pooling, several other pooling strategies have been explored. Max pooling takes the maximum value in each dimension across all token vectors, which can emphasize strong signals from individual tokens. Weighted mean pooling assigns importance weights to tokens before averaging, potentially downweighting stopwords and upweighting content words. Some specialized architectures produce multiple vectors per document (one per sentence or paragraph) to support more granular retrieval.

The Sentence Transformers library, which we use in our code implementation, uses mean pooling combined with optional L2 normalization by default. This combination has proven reliable across many retrieval benchmarks. The normalization step ensures that embedding magnitudes are uniform, which simplifies the similarity computation and prevents magnitude biases from distorting retrieval results.

One important consideration is the maximum sequence length. Transformer models have a fixed context window, typically 512 tokens for BERT-based models. Documents longer than this limit must be truncated or split into multiple chunks. In RAG systems, documents are usually pre-chunked into passages of a few hundred words specifically to fit within these limits. Retrieval quality depends on both chunk size and overlap, including whether chunks respect sentence boundaries. Longer chunks preserve more context but may dilute relevance signals; shorter chunks are more focused but may miss context needed to answer the query.