BERT Representations › Pooling Strategies: Beyond CLS
Language AI HandbookAll chapters BERT Representations › Pooling Strategies: Beyond CLSOne section · Sections in this chapter4 / 10 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
The [CLS] token is one way to get a sentence representation, but it's not the only way. Pooling strategies aggregate information across all token positions.
The underlying question pooling answers is: given that every token in the sequence has a rich contextual representation, how should we combine them into a single fixed-size vector? Each token's representation already encodes context from the entire sequence through self-attention, so the representations aren't independent. A pooled representation tries to summarize this distributed information into something compact and task-relevant.
Notice that the [CLS] approach is itself a pooling strategy of sorts, one where all the aggregation work happens inside the transformer through attention, and only the special token's output is used. The alternative pooling methods described here perform the aggregation after the transformer, using simple arithmetic over the output token representations. Both approaches have their place, and the best choice depends on what you want the representation to capture.
In practice, when building systems that need to compare sentence meanings, such as question-answering retrievers or document similarity engines, mean pooling over the final layer often outperforms [CLS] without any additional fine-tuning. This is a quick win: simply change how you read the output tensor, and your similarity quality improves. The reason is that mean pooling spreads the representation load across all content tokens rather than concentrating it at a single position that was not specifically trained for generic similarity.
Mean Pooling
Mean pooling averages the representations of all tokens (typically excluding special tokens).
def mean_pooling(hidden_state, attention_mask):
"""
Average token representations, respecting the attention mask.
Args:
hidden_state: [batch, seq_len, hidden_dim]
attention_mask: [batch, seq_len], 1 for real tokens, 0 for padding
"""
# Expand mask to hidden dimension
mask_expanded = attention_mask.unsqueeze(-1).float()
# Sum of non-masked representations
sum_embeddings = torch.sum(hidden_state * mask_expanded, dim=1)
# Count of non-masked tokens
sum_mask = mask_expanded.sum(dim=1).clamp(min=1e-9)
return sum_embeddings / sum_mask
# Apply to our example
last_hidden = outputs.last_hidden_state # [1, seq_len, 768]
attention_mask = inputs["attention_mask"] # [1, seq_len]
mean_pooled = mean_pooling(last_hidden, attention_mask)Mean Pooling Result ---------------------------------------- Input shape: torch.Size([1, 9, 768]) Output shape: torch.Size([1, 768]) Mean pooled norm: 9.2480 CLS norm: 14.8434
Mean pooling gives equal weight to every token. This works well when all parts of the sentence contribute equally to the overall meaning, which is often the case for similarity tasks. The attention mask is essential here: without it, padding tokens would dilute the representation by contributing zero-content vectors. Always pass the mask when pooling to ensure that only real tokens contribute to the average.
Max Pooling
Max pooling takes the maximum value across the sequence for each dimension.
def max_pooling(hidden_state, attention_mask):
"""
Take max over token positions for each dimension.
Padding tokens are set to large negative values to exclude them.
"""
# Expand mask to match hidden state dimensions
mask_expanded = attention_mask.unsqueeze(-1).expand_as(hidden_state).float()
# Replace padding positions with very negative values
hidden_state = hidden_state.masked_fill(mask_expanded == 0, -1e9)
# Max over sequence dimension
max_pooled, _ = torch.max(hidden_state, dim=1)
return max_pooled
max_pooled = max_pooling(last_hidden, attention_mask)Max Pooling Result ---------------------------------------- Output shape: torch.Size([1, 768]) Max pooled norm: 18.1580 Values range: [-0.5583, 3.6662]
Max pooling captures the strongest signals in each dimension. It can be useful when specific keywords or phrases are most important, though it tends to be more sensitive to outliers than mean pooling. The intuition is that for some tasks, presence matters more than prevalence. If a document mentions "cancer" once in 500 tokens, mean pooling dilutes that signal, while max pooling preserves the highest activation in the "cancer-related" dimensions regardless of how rarely the topic appeared. For topic labeling and keyword-sensitive retrieval, this property can be an advantage.
Comparing Pooling Strategies
Let's compare how different pooling strategies handle the same sentences.
def get_all_poolings(text, model, tokenizer):
"""Extract CLS, mean, and max pooled representations."""
inputs = tokenizer(text, return_tensors="pt", truncation=True)
with torch.no_grad():
outputs = model(**inputs)
hidden = outputs.last_hidden_state
mask = inputs["attention_mask"]
return {
"cls": hidden[0, 0, :].numpy(),
"mean": mean_pooling(hidden, mask).squeeze().numpy(),
"max": max_pooling(hidden, mask).squeeze().numpy(),
}
# Test on similar and dissimilar sentence pairs
test_pairs = [
(
"A dog is running through the park.",
"A canine sprints across the garden.",
), # Similar
(
"A dog is running through the park.",
"Stock prices fell sharply today.",
), # Dissimilar
]
pooling_results = []
for s1, s2 in test_pairs:
p1 = get_all_poolings(s1, model, tokenizer)
p2 = get_all_poolings(s2, model, tokenizer)
result = {"pair": (s1[:35], s2[:35])}
for method in ["cls", "mean", "max"]:
sim = cosine_similarity([p1[method]], [p2[method]])[0, 0]
result[method] = sim
pooling_results.append(result)Pooling Strategy Comparison ====================================================================== Similar pair: 'A dog is running through the park....' 'A canine sprints across the garden....' CLS similarity: 0.9445 Mean similarity: 0.8737 Max similarity: 0.9312 Dissimilar pair: 'A dog is running through the park....' 'Stock prices fell sharply today....' CLS similarity: 0.8200 Mean similarity: 0.5439 Max similarity: 0.7955
All strategies show the similar pair as more similar than the dissimilar pair, but the margins differ. Mean pooling often produces better separation for similarity tasks because it incorporates information from all content tokens rather than relying solely on the [CLS] position. The discrimination gap, meaning how much larger the similarity is for the matching pair than for the unrelated pair, is the number that matters most here. A pooling method with a tiny gap will give you low-confidence rankings in retrieval systems, making it hard to separate relevant from irrelevant results at query time.
Attention-Weighted Pooling
A more sophisticated approach uses the model's own attention weights to determine token importance.
def attention_weighted_pooling(hidden_state, attention_weights):
"""
Pool using attention weights as importance scores.
Args:
hidden_state: [batch, seq_len, hidden_dim]
attention_weights: [batch, num_heads, seq_len, seq_len]
"""
# Average attention across heads and queries to get per-token importance
# Shape: [batch, seq_len]
token_importance = attention_weights.mean(dim=(1, 2))
# Normalize to sum to 1
token_importance = token_importance / token_importance.sum(
dim=-1, keepdim=True
)
# Weighted average
weighted = torch.einsum("bs,bsh->bh", token_importance, hidden_state)
return weighted
# Use eager attention backend so output_attentions=True returns actual weights
# (SDPA/Flash Attention backends do not compute or return attention matrices)
model_eager = BertModel.from_pretrained(
"bert-base-uncased", attn_implementation="eager"
)
model_eager.train(False)
# Get attention weights
with torch.no_grad():
outputs_with_attn = model_eager(**inputs, output_attentions=True)
# Use last layer's attention
last_attn = outputs_with_attn.attentions[-1] # [1, 12, seq_len, seq_len]
attn_pooled = attention_weighted_pooling(
outputs_with_attn.last_hidden_state, last_attn
)Attention-Weighted Pooling ---------------------------------------- Attention shape: torch.Size([1, 12, 9, 9]) Pooled output shape: torch.Size([1, 768]) Output norm: 10.8894
Attention-weighted pooling lets the model decide which tokens matter most. Tokens that receive more attention across the sequence get higher weight in the final representation. In this example, content-bearing words like "bank", "approved", and "loan" naturally receive higher importance scores than function words like "the".
The appeal of attention-weighted pooling is that it's data-driven rather than heuristic. Rather than uniformly averaging or blindly taking the maximum, the weights come from the model's own learned attention patterns. Notice, however, that this approach has a subtle limitation: attention weights reflect what the model found important for its pre-training objective, which may not align perfectly with what's important for your task. A model trained on next-sentence prediction might attend heavily to discourse markers and topic words, but for a technical document retrieval task, precise terminology might matter more. In practice, attention-weighted pooling adds engineering complexity without reliably outperforming mean pooling on most benchmarks, so treat it as a specialized tool rather than a default.