Embedding Models: Architecture, Pooling & Selection › Pooling Strategies

The pooling layer is arguably the most necessary design choice in an embedding model. It determines how the rich, token-level information from the transformer encoder is compressed into a single vector. To understand why this choice matters so much, consider what happens at the output of a transformer encoder. You have a matrix of hidden states, one row per token, each row a high-dimensional vector capturing that token's meaning in context. The pooling layer must somehow condense this entire matrix, which can have hundreds of rows for longer inputs, into a single row. Different strategies for performing this condensation emphasize different aspects of the input, and the choice of strategy interacts deeply with how the model was trained.

Think of the pooling operation as deciding how to summarize a meeting transcript. You could take the notes from the meeting chair (CLS pooling), who was designated as the official recorder but might have focused on their own agenda items. You could average everyone's notes together (mean pooling), which captures the collective sense of the meeting but might dilute points that were only raised by one person. You could take the most emphatic point from each person (max pooling), which highlights the strongest sentiments but misses the fine-grained discussion. Or you could take the notes from the last person to speak (last token pooling), who had heard the entire discussion and could synthesize it, but might have overweighted the final points.

BERT and its variants include a special [CLS] token at the beginning of every input sequence. During BERT's pre-training, the hidden state corresponding to [CLS] is used as the "aggregate sequence representation" for the next sentence prediction task: a linear classifier applied to this single vector must decide whether the second sentence follows the first. CLS pooling simply takes this one vector as the sentence embedding, designating the first position as the dedicated aggregation point.

Formally, the CLS pooling operation extracts the hidden state at the very first position of the output:

$$\mathbf{e} = \mathbf{h}_{\text{[CLS]}}$$

where:

• $\mathbf{e}$: the resulting sentence embedding vector of dimension $d$
• $\mathbf{h}_{\text{[CLS]}}$: the final hidden state vector at the [CLS] token position, which is always the first token in BERT's tokenization scheme

The appeal of CLS pooling is that the model has a dedicated token whose job is to aggregate information from the entire sequence. Through bidirectional self-attention, the [CLS] token can attend to every other token at every layer, and the model can learn to pack a summary of the input into this position. Because self-attention is fully connected (every token can interact with every other token), the [CLS] token has, in principle, access to the full content of the input by the time it reaches the final layer. If the training objective rewards the [CLS] position for capturing the overall meaning of the input, the model has both the mechanism (self-attention) and the incentive (the loss function) to make this work.

However, CLS pooling has a significant weakness: the [CLS] token's representation is heavily shaped by its pre-training objective (next sentence prediction in BERT), which does not align well with semantic similarity. The next sentence prediction task is a binary classification problem. It asks a coarse question: does sentence B follow sentence A in the original document? This coarse signal does not require, and does not reward, the model for encoding fine-grained semantic meaning into the [CLS] representation. As a result, without explicit fine-tuning for embedding quality, [CLS] representations can be poorly organized for semantic search, placing dissimilar texts near each other and similar texts far apart. After proper contrastive fine-tuning that specifically rewards the [CLS] representation for semantic discrimination, CLS pooling works well and is used by several strong models including DeBERTa-based models and many decoder-based embedding models.

Mean pooling takes a fundamentally different approach from CLS pooling. Rather than relying on a single designated token to carry all the information, mean pooling distributes the responsibility across every token in the sequence. It computes the sentence embedding by averaging the hidden states of all input tokens, excluding padding tokens, to produce a single representative vector.

To understand the motivation, consider what happens at the final layer of a transformer encoder. Each token's hidden state is a rich, contextualized representation that has been refined through many layers of self-attention and feed-forward processing. The hidden state at the word "photosynthesis" in a biology passage does not just encode the word itself: it encodes the word in the context of the entire passage, informed by every other word through self-attention. By averaging all of these context-rich representations, mean pooling creates a vector that reflects the collective semantic content of the entire sequence, weighted equally across all positions. The formal expression is:

$$\mathbf{e} = \frac{1}{|\mathcal{T}|} \sum_{i \in \mathcal{T}} \mathbf{h}_i$$

where:

• $\mathbf{e}$: the resulting sentence embedding vector
• $\mathcal{T}$: the set of indices corresponding to real (non-padding) tokens
• $\mathbf{h}_i$: the hidden state vector output by the transformer at position $i$
• $|\mathcal{T}|$: the count of real tokens in the sequence, which is the normalizing denominator

In practice, this is implemented using the attention mask to zero out padding positions before averaging. Padding tokens are added to make all sequences in a batch have the same length, but they carry no meaningful content and should not influence the embedding. The attention mask, a binary vector with 1 for real tokens and 0 for padding, provides exactly the information needed to exclude padding positions:

$$\mathbf{e} = \frac{\sum_{i=1}^{n} m_i \cdot \mathbf{h}_i}{\sum_{i=1}^{n} m_i}$$

where:

• $\mathbf{e}$: the resulting sentence embedding vector
• $n$: the total length of the tokenized sequence, including padding
• $m_i$: the attention mask value at position $i$ (1 for real tokens, 0 for padding)
• $\mathbf{h}_i$: the hidden state vector at position $i$

The numerator sums only the hidden states of real tokens, since multiplying by the mask zeroes out padding positions. The denominator counts the number of real tokens. This ensures that the average is taken only over meaningful positions, regardless of how much padding was added.

Mean pooling is the most widely used strategy in modern embedding models, and for good reason. By averaging over all tokens, it captures information distributed across the entire sequence rather than relying on a single position that might or might not have attended equally to all parts of the input. It is more reliable to the arbitrary characteristics of specific token positions, and the training signal can flow to every token equally, giving the model more flexibility in how it distributes semantic information across the sequence. For longer inputs, where useful information may appear anywhere from the opening sentence to the final clause, this distributed representation is especially valuable.

Sentence-BERT showed that mean pooling consistently outperforms CLS pooling when fine-tuning from a pre-trained BERT checkpoint. Most state-of-the-art models, including the E5, GTE, and BGE families, use mean pooling as their default strategy. The reason mean pooling often outperforms CLS pooling is not that CLS is inherently inferior, but that mean pooling is more forgiving: even a model that has not perfectly learned to concentrate information at the CLS position will produce reasonable mean-pooled embeddings, because the average captures information wherever it happens to reside in the hidden states.

A refinement of mean pooling introduces position-dependent weights, letting the model to emphasize certain token positions over others. Typically, later tokens in the sequence receive higher weight, based on the intuition that they have "seen" more context through the attention mechanism and may carry more refined semantic information:

$$\mathbf{e} = \frac{\sum_{i=1}^{n} w_i \cdot m_i \cdot \mathbf{h}_i}{\sum_{i=1}^{n} w_i \cdot m_i}$$

where:

• $\mathbf{e}$: the resulting sentence embedding vector
• $n$: the total length of the tokenized sequence
• $w_i$: the weight assigned to position $i$ (common choices include $w_i = i$ for linear weighting or $w_i = i^2$ for quadratic weighting)
• $m_i$: the attention mask value at position $i$
• $\mathbf{h}_i$: the hidden state vector at position $i$

The intuition for this weighting is that in some models or some types of text, information is not uniformly distributed across positions. For instance, the predicate of a sentence (often appearing in the middle or near the end) might be more semantically informative than the subject (often appearing near the beginning). In a bidirectional encoder where every token attends to every other token, this intuition is weaker than it might seem, because the model can freely route information to any position. However, in practice, positional weighting can provide a small boost in some settings.

In production models, weighted mean pooling offers marginal improvements and is not widely adopted. The added complexity of choosing and justifying a weighting scheme rarely justifies the modest gains, and standard mean pooling remains the preferred default. We mention it here because you may encounter it in research papers or as a hyperparameter option in some libraries.

Max pooling takes the element-wise maximum across all token positions, selecting the strongest activation for each dimension independently:

$$e_j = \max_{i \in \mathcal{T}} h_{i,j}$$

where:

• $e_j$: the value of the $j$-th dimension in the final embedding vector
• $\mathcal{T}$: the set of non-padding token positions
• $h_{i,j}$: the scalar value of the $j$-th dimension of the hidden state at token position $i$

To understand what this formula does, consider a single dimension $j$ of the embedding. Across all token positions in the input, each token's hidden state contributes a scalar value for this dimension. Max pooling selects the largest of these values. The resulting embedding vector is assembled dimension by dimension, where each dimension's value comes from whichever token produced the strongest activation along that particular feature axis. Different dimensions can be satisfied by different tokens: dimension 47 might take its value from token position 3, while dimension 312 takes its value from token position 11.

Max pooling captures the presence of a feature anywhere in the sequence, regardless of how infrequently it appears. If a particular dimension activates strongly when the model detects a mention of a geographic location, max pooling will preserve that signal even if the location is mentioned only once in a long passage. Mean pooling would dilute this signal by averaging with all the other tokens where that dimension has lower activation. In image processing, max pooling over spatial regions is a standard technique for detecting the presence of local features regardless of position, and the intuition is similar for text.

However, max pooling is sensitive to outlier activations: a single anomalous token can dominate the embedding along many dimensions simultaneously, creating a vector that reflects the most extreme aspects of the input rather than its typical character. It is also more sensitive to padding-related artifacts, which is why the mask-based exclusion of padding tokens is particularly important. Max pooling is rarely used as the sole pooling strategy in modern production models, but it is sometimes combined with mean pooling in ensemble approaches.

For decoder-only models like GPT or LLaMA that use causal attention, the information flow through the transformer has a fundamentally different directional character than in bidirectional encoders. In causal attention, each token can only attend to itself and to tokens that precede it in the sequence. The first token sees only itself; the second token sees the first and itself; and each successive token has access to a progressively larger window of context. By the time we reach the last token, it has attended to every previous token in the sequence, which makes it the only position that has incorporated information from the complete input:

$$\mathbf{e} = \mathbf{h}_n$$

where:

• $\mathbf{e}$: the resulting sentence embedding vector
• $\mathbf{h}_n$: the hidden state vector at the last non-padding token position $n$

This is the analog of CLS pooling for causal models, but the reasoning for using the last position rather than the first is almost exactly reversed. In a bidirectional encoder, the [CLS] token at the beginning can attend to all other tokens through bidirectional self-attention, which makes it a well-informed aggregation point. In a causal decoder, the same logic applies to the last token: it is the last position, not the first, that has access to the full sequence. Using the first token for a causal model would be deeply counterproductive, because the first token has attended only to itself and has no information about what follows.

Models like E5-Mistral, SFR-Embedding, and GTE-Qwen2 use this strategy. Some models append a special end-of-sequence token or a dedicated aggregation token (such as </s> or [EOS]) to each input and use its representation. This special token acts as a dedicated aggregation point, analogous to BERT's [CLS], but positioned at the end where it can attend to the full prefix. Adding a specific token for this purpose gives the model a consistent, well-defined position to learn to concentrate embedding information, rather than relying on the last token of the input text, which varies in identity across different inputs.

The practical consequence of last token pooling is that decoder-based embedding models often benefit from appending a short instruction or query prefix to the input text, because the final position then reflects the task context as well as the text content. This is why models like E5-Mistral expect inputs formatted as "Instruct: Retrieve relevant passages. Query: {query text}" or similar patterns.

The choice of pooling strategy interacts strongly with the model architecture and training procedure, and this interaction is one of the most important things to understand about embedding models. A pooling strategy cannot be evaluated in isolation, because its effectiveness depends entirely on how the model was designed and trained to use it. A CLS token from a model trained with mean pooling is not the same kind of object as a CLS token from a model trained to use CLS pooling: in the former case, the training never rewarded the CLS position for putting semantic information, so it may contain arbitrary or residual information rather than a meaningful summary.

For encoder models (BERT, RoBERTa, DeBERTa), mean pooling generally wins over CLS pooling when fine-tuned with contrastive objectives, because the training signal distributes more evenly across token positions. CLS pooling is competitive after proper fine-tuning where the loss specifically rewards the CLS representation for semantic discrimination. For decoder-only models (Mistral, LLaMA, Qwen), last token pooling is the natural and appropriate choice due to causal attention masking: it is the only position with sufficient context to serve as a standalone summary. For encoder-decoder models (T5, BART), the encoder's output is typically mean-pooled, using the same reasoning that favors mean pooling in encoder-only models, since the encoder in these architectures uses bidirectional attention.

The practical guideline is simple: always use the pooling strategy specified in the model's documentation or that was used during its training. Mixing pooling strategies produces degraded embeddings not because one strategy is universally better, but because the model's weights encode expectations about how the pooling layer will read out its hidden states. Violating those expectations corrupts the information being read out.

Let's implement the main pooling strategies and see how they differ in practice. We'll encode three sentences through the same model and apply each pooling method to the raw hidden states, then compare the resulting cosine similarity matrices.

The heatmaps reveal the differences between pooling strategies. This particular model (all-MiniLM-L6-v2) was trained with mean pooling, so that strategy produces the most semantically meaningful similarities: the two animal-on-surface sentences are highly similar, while the stock market sentence is more distant. CLS and max pooling, applied to a model trained for mean pooling, produce less differentiated similarity scores. The practical point is that the pooling strategy must match what the model was trained with. Using CLS pooling on a model fine-tuned with mean pooling, or vice versa, will degrade quality in ways that may not be immediately obvious but will show up as poorer retrieval performance in your RAG pipeline.

The key parameters for the tokenizer used in this example are:

• padding: Adds special padding tokens to ensure all sequences in a batch have the same length, letting batch matrix operations.
• truncation: Truncates sequences that exceed the model's maximum context length to prevent out-of-range errors.
• return_tensors: Specifies the return type ("pt" for PyTorch) for compatibility with the model.