Pooling Strategies: From Multiple Tokens to One Vector

Capsule overview

After the Transformer processes the tokens, you have multiple hidden states (one vector per token). To obtain a full-sentence embedding, you need "pooling": reducing multiple vectors to a single one.

In this capsule you'll learn the 3 main pooling strategies (mean, CLS, max), implement each one with numpy, compare their trade-offs, and understand why mean pooling dominates in sentence embeddings. You'll also see practical code that simulates the pooling process.

By the end, you'll be able to make informed decisions about which pooling strategy to use based on your use case.

The pooling problem

Situation:

# After the Transformer encoder:
text = "Python is a language"

# Hidden states (one vector per token):
hidden_states = [
    [0.8, 0.1, 0.3, 0.5, 0.2],  # "Python"
    [0.9, 0.3, 0.2, 0.4, 0.1],  # "is"
    [0.7, 0.2, 0.4, 0.6, 0.3],  # "a"
    [0.6, 0.4, 0.1, 0.3, 0.4]   # "language"
]

# You need 1 single embedding for the whole sentence:
sentence_embedding = ???  # How do we get it?

Solution: Pooling (aggregation).

Strategy #1: Mean Pooling

What it is:

The average of all the hidden states (token vectors).

Formula:

embedding = (v₁ + v₂ + v₃ + ... + vₙ) / n

Where:
- v₁, v₂, ... = The hidden states of each token
- n = The number of tokens

Implementation with numpy:

import numpy as np

def mean_pooling(hidden_states):
    """
    Mean pooling: The average of all hidden states
    
    Args:
        hidden_states: Array of shape (n_tokens, hidden_size)
    
    Returns:
        An embedding of shape (hidden_size,)
    """
    return np.mean(hidden_states, axis=0)

# Example
hidden_states = np.array([
    [0.8, 0.1, 0.3, 0.5, 0.2],  # Token 1
    [0.9, 0.3, 0.2, 0.4, 0.1],  # Token 2
    [0.7, 0.2, 0.4, 0.6, 0.3],  # Token 3
    [0.6, 0.4, 0.1, 0.3, 0.4]   # Token 4
])

embedding = mean_pooling(hidden_states)
print(f"Embedding: {embedding}")
print(f"Shape: {embedding.shape}")

Output:

Embedding: [0.75 0.25 0.25 0.45 0.25]
Shape: (5,)

Explanation:

• Dim 0: (0.8 + 0.9 + 0.7 + 0.6) / 4 = 0.75
• Dim 1: (0.1 + 0.3 + 0.2 + 0.4) / 4 = 0.25
• ...

Advantages of mean pooling:

✅ 1. Simple and effective

# One line of code:
embedding = np.mean(hidden_states, axis=0)

✅ 2. Considers all tokens

# Each token contributes to the final embedding
# You don't discard information

✅ 3. Robust to variable length

# Works the same for 5 tokens or 500 tokens
embedding_short = mean_pooling(hidden_states_5)    # Shape: (768,)
embedding_long = mean_pooling(hidden_states_500)   # Shape: (768,)

Disadvantages of mean pooling:

❌ 1. Irrelevant tokens dilute

text = "Python, um, is, like, a language"
# "um", "like" contribute to the average even though they're noise

# Mean pooling: Averages EVERYTHING (including noise)

❌ 2. Doesn't weight importance

# "Python" and "is" contribute equally to the average
# But "Python" is more semantically important

Strategy #2: CLS Pooling

What it is:

Use only the hidden state of the [CLS] token (a special token added at the start).

BERT adds [CLS] at the start:

text = "Python is a language"

# BERT tokenization:
tokens = ["[CLS]", "Python", "is", "a", "language"]

# Hidden states:
hidden_states = [
    [0.5, 0.6, 0.7, 0.8, 0.9],  # [CLS] ← We use ONLY this one
    [0.8, 0.1, 0.3, 0.5, 0.2],  # Python
    [0.9, 0.3, 0.2, 0.4, 0.1],  # is
    [0.7, 0.2, 0.4, 0.6, 0.3],  # a
    [0.6, 0.4, 0.1, 0.3, 0.4]   # language
]

# CLS pooling: Take only hidden_states[0]
embedding = hidden_states[0]  # [0.5, 0.6, 0.7, 0.8, 0.9]

Implementation with numpy:

def cls_pooling(hidden_states):
    """
    CLS pooling: Use only the first hidden state ([CLS] token)
    
    Args:
        hidden_states: Array of shape (n_tokens, hidden_size)
    
    Returns:
        An embedding of shape (hidden_size,)
    """
    return hidden_states[0]

# Example (assuming [CLS] is the first token)
hidden_states = np.array([
    [0.5, 0.6, 0.7, 0.8, 0.9],  # [CLS]
    [0.8, 0.1, 0.3, 0.5, 0.2],  # Token 1
    [0.9, 0.3, 0.2, 0.4, 0.1],  # Token 2
])

embedding = cls_pooling(hidden_states)
print(f"Embedding: {embedding}")

Output:

Embedding: [0.5 0.6 0.7 0.8 0.9]

Advantages of CLS pooling:

✅ 1. A specifically trained token

# [CLS] is trained to "summarize" the whole sentence
# During BERT's training, [CLS] learns to capture the global meaning

✅ 2. Doesn't dilute with irrelevant tokens

# Noise tokens don't affect it (they're not averaged)

Disadvantages of CLS pooling:

❌ 1. Ignores other tokens

# Only uses [CLS], discards information from tokens 2-N
# Potential loss of nuance

❌ 2. Requires a special token

# OpenAI embeddings do NOT use [CLS]
# Only works with BERT-style models

Strategy #3: Max Pooling

What it is:

Take the maximum value of each dimension across all tokens.

def max_pooling(hidden_states):
    """
    Max pooling: The maximum per dimension
    
    Args:
        hidden_states: Array of shape (n_tokens, hidden_size)
    
    Returns:
        An embedding of shape (hidden_size,)
    """
    return np.max(hidden_states, axis=0)

# Example
hidden_states = np.array([
    [0.8, 0.1, 0.3, 0.5, 0.2],
    [0.9, 0.3, 0.2, 0.4, 0.1],
    [0.7, 0.2, 0.4, 0.6, 0.3],
    [0.6, 0.4, 0.1, 0.3, 0.4]
])

embedding = max_pooling(hidden_states)
print(f"Embedding: {embedding}")

Output:

Embedding: [0.9 0.4 0.4 0.6 0.4]

Explanation:

• Dim 0: max(0.8, 0.9, 0.7, 0.6) = 0.9
• Dim 1: max(0.1, 0.3, 0.2, 0.4) = 0.4
• ...

Advantages of max pooling:

✅ Captures the most "salient" features

# If a token has a very high value in one dimension,
# that value is preserved in the final embedding

Disadvantages of max pooling:

❌ 1. Loses average information

# It only captures the maximum values, not the "general sense"

❌ 2. Less used in practice

# Mean pooling and CLS pooling dominate
# Max pooling is rare in sentence embeddings

Practical comparison: 3 pooling strategies

Full comparison code:

import numpy as np
from openai import OpenAI
import os
from dotenv import load_dotenv

# Setup
load_dotenv()
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# Simulate hidden states (in production they come from the Transformer)
# For demonstration, we generate random ones
np.random.seed(42)
hidden_states = np.random.randn(10, 5)  # 10 tokens, 5 dims

print("Hidden States (10 tokens × 5 dims):")
print(hidden_states)
print()

# Mean pooling
mean_emb = np.mean(hidden_states, axis=0)
print(f"Mean Pooling: {mean_emb}")

# CLS pooling (take the first token)
cls_emb = hidden_states[0]
print(f"CLS Pooling:  {cls_emb}")

# Max pooling
max_emb = np.max(hidden_states, axis=0)
print(f"Max Pooling:  {max_emb}")

# Compare with cosine similarity
def cosine_similarity(vec_a, vec_b):
    return np.dot(vec_a, vec_b) / (np.linalg.norm(vec_a) * np.linalg.norm(vec_b))

print(f"\nComparison:")
print(f"Mean vs CLS: {cosine_similarity(mean_emb, cls_emb):.4f}")
print(f"Mean vs Max: {cosine_similarity(mean_emb, max_emb):.4f}")
print(f"CLS vs Max:  {cosine_similarity(cls_emb, max_emb):.4f}")

Output (example):

Mean Pooling: [ 0.12 -0.05  0.18 -0.08  0.15]
CLS Pooling:  [ 0.50  0.60  0.70  0.80  0.90]
Max Pooling:  [ 1.45  0.85  1.20  0.95  1.10]

Comparison:
Mean vs CLS: 0.65  ← Somewhat similar
Mean vs Max: 0.78  ← More similar (mean smooths max)
CLS vs Max:  0.55  ← Less similar

Which pooling each model uses

Reference table:

Trend: Mean pooling dominates in sentence embeddings (SBERT popularized it).

Attention Mask: Correct pooling with padding

Problem: Sequences of different length

# Batch of 3 sentences:
texts = [
    "Python",                    # 1 token
    "Python is a language",      # 4 tokens
    "JavaScript"                 # 1 token
]

# To process in a batch, you need the SAME size:
# Padding (fill with special tokens):
# [Python, PAD, PAD, PAD]
# [Python, is, a, language]
# [JavaScript, PAD, PAD, PAD]

Mean pooling with attention mask:

def mean_pooling_with_mask(hidden_states, attention_mask):
    """
    Mean pooling ignoring padding tokens
    
    Args:
        hidden_states: (batch_size, seq_len, hidden_size)
        attention_mask: (batch_size, seq_len) - 1 = real token, 0 = padding
    
    Returns:
        Embeddings: (batch_size, hidden_size)
    """
    # Expand the attention mask for multiplication
    attention_mask_expanded = np.expand_dims(attention_mask, axis=-1)
    
    # Multiply hidden states by the mask (zeroes out padding)
    masked_hidden_states = hidden_states * attention_mask_expanded
    
    # Sum (only real tokens)
    sum_hidden = np.sum(masked_hidden_states, axis=1)
    
    # Divide by the number of real tokens
    sum_mask = np.sum(attention_mask_expanded, axis=1)
    sum_mask = np.clip(sum_mask, a_min=1e-9, a_max=None)  # Avoid division by 0
    
    # Mean
    embeddings = sum_hidden / sum_mask
    
    return embeddings

# Example
hidden_states = np.array([
    [[0.8, 0.1, 0.3], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],  # Batch 1: 1 real token
    [[0.9, 0.3, 0.2], [0.7, 0.2, 0.4], [0.6, 0.4, 0.1]]   # Batch 2: 3 real tokens
])

attention_mask = np.array([
    [1, 0, 0],  # Batch 1: Only the first token is real
    [1, 1, 1]   # Batch 2: All tokens are real
])

embeddings = mean_pooling_with_mask(hidden_states, attention_mask)
print("Embeddings:")
print(embeddings)

Output:

Embeddings:
[[0.8  0.1  0.3 ]  ← Batch 1: Only used token 1
 [0.73 0.3  0.23]]  ← Batch 2: Average of 3 tokens

Empirical comparison: Mean vs CLS

Experiment: Which is better?

Setup: Evaluate both pooling strategies on a similarity task.

from sentence_transformers import SentenceTransformer

# Load the SBERT model (uses mean pooling by default)
model = SentenceTransformer('all-MiniLM-L6-v2')

# Pairs of similar sentences
pairs = [
    ("The cat sleeps", "The feline rests"),
    ("Python is popular", "Python is widely used"),
    ("I like pizza", "I love pizza")
]

# Generate embeddings with mean pooling (default)
for text_a, text_b in pairs:
    emb_a = model.encode(text_a)
    emb_b = model.encode(text_b)
    
    sim = np.dot(emb_a, emb_b) / (np.linalg.norm(emb_a) * np.linalg.norm(emb_b))
    print(f"'{text_a}' vs '{text_b}': {sim:.4f}")

Output with mean pooling:

'The cat sleeps' vs 'The feline rests': 0.78
'Python is popular' vs 'Python is widely used': 0.85
'I like pizza' vs 'I love pizza': 0.82

Mean pooling captures paraphrases effectively.

Exercises

Exercise 1: Implement mean pooling

Implement mean pooling for these hidden states:

import numpy as np

hidden_states = np.array([
    [1.0, 2.0, 3.0],
    [4.0, 5.0, 6.0],
    [7.0, 8.0, 9.0]
])

# Calculate mean pooling

def mean_pooling(hidden_states):
    return np.mean(hidden_states, axis=0)

embedding = mean_pooling(hidden_states)
print(f"Embedding: {embedding}")

Embedding: [4. 5. 6.]

Exercise 2: CLS vs Mean pooling

Calculate CLS and mean pooling for the same hidden states. Are they similar?

hidden_states = np.array([
    [0.5, 0.6, 0.7],  # [CLS]
    [0.8, 0.1, 0.3],
    [0.9, 0.3, 0.2]
])

# Calculate both and compare

# CLS pooling
cls_emb = hidden_states[0]

# Mean pooling
mean_emb = np.mean(hidden_states, axis=0)

print(f"CLS:  {cls_emb}")
print(f"Mean: {mean_emb}")

# Compare
similarity = np.dot(cls_emb, mean_emb) / (np.linalg.norm(cls_emb) * np.linalg.norm(mean_emb))
print(f"\nSimilarity: {similarity:.4f}")

CLS:  [0.5 0.6 0.7]
Mean: [0.73 0.33 0.4]

Similarity: 0.9850

Exercise 3: Max pooling

Implement max pooling:

hidden_states = np.array([
    [0.8, 0.1, 0.3],
    [0.9, 0.3, 0.2],
    [0.7, 0.2, 0.4]
])

# Implement max pooling

def max_pooling(hidden_states):
    return np.max(hidden_states, axis=0)

embedding = max_pooling(hidden_states)
print(f"Embedding: {embedding}")

Embedding: [0.9 0.3 0.4]

Summary

What you learned:

• ✅ Pooling: Reduce multiple tokens → 1 embedding
• ✅ Mean pooling: Average (all tokens matter)
• ✅ CLS pooling: The special [CLS] token (BERT-style)
• ✅ Max pooling: Maximum per dimension (rare)
• ✅ Attention mask: Ignore padding in batch processing
• ✅ Best practice: Mean pooling for sentence embeddings

Key concepts:

1. Mean pooling dominates in modern sentence embeddings
2. CLS pooling requires a special token (BERT)
3. The attention mask is critical for batch processing

Additional resources

1. Sentence-BERT Paper - Justifies mean pooling
2. Pooling Strategies Comparison - Experiments
3. BERT Pooling - CLS tutorial
4. Attention Mask Explained - HuggingFace
5. Mean Pooling Implementation - SBERT code

In the next capsule

Capsule 06: Normalization

You'll learn:

• L2 normalization (numpy code)
• Why normalize embeddings
• Dot product vs cosine similarity
• When to normalize (depends on the model)
• Practical implementation

From pooling to normalization.

Module 2 - Embeddings Deep Dive Guide Aggregating tokens: from many vectors to one