Pool Token Embeddings into Sequence Representations

Transformer encoders usually produce one vector for every input token. Many downstream tasks, however, need one vector for the whole input: a classifier may need a single representation of a support ticket, and a retrieval system may need one vector for an entire passage.

The step that converts a variable number of token vectors into one fixed-size vector is pooling. It looks simple, but the choice of pooling rule changes what information survives, how padding must be handled, and whether the resulting vector matches the way a model was trained.

This article builds a practical mental model for sequence pooling, starts with masked mean pooling, and then explains when special-token, max, or learned pooling is a better fit.

Pooling changes the shape, not just the size

Suppose an encoder maps a five-token input to five vectors, each with four dimensions:

input tokens:       5
embedding dimension: 4
encoder output:      5 x 4

A downstream classifier that expects one four-dimensional vector cannot consume that matrix directly. Pooling reduces the token dimension:

5 x 4 token matrix -> pooling -> 1 x 4 sequence vector

The important question is what the reduction means. Averaging says every included token contributes equally. Max pooling keeps the strongest value in each dimension. Selecting a designated token assumes the model has learned to store useful sequence-level information at that position.

These operations have the same output shape but different semantics.

Start with masked mean pooling

Mean pooling is a useful baseline because it has a clear interpretation and no learned pooling parameters. For token vectors h_1 ... h_n, the unmasked mean is:

pooled = (h_1 + h_2 + ... + h_n) / n

Consider three two-dimensional token vectors:

h_1 = [1.0, 0.0]
h_2 = [0.5, 1.0]
h_3 = [0.0, 2.0]

Their mean is:

[(1.0 + 0.5 + 0.0) / 3,
 (0.0 + 1.0 + 2.0) / 3]
= [0.5, 1.0]

The result summarizes the included token representations component by component. It does not reconstruct the original sequence, and it discards token order at the pooling step. Order can still influence each token vector because the encoder computed those vectors in context before pooling.

Padding must not enter the average

Batched sequences commonly have different lengths. Shorter examples are padded so tensors share a rectangular shape:

example A: [t1, t2, t3, t4]
example B: [t1, t2, PAD, PAD]

A plain average over all four positions gives example B’s padding positions the same weight as real tokens. Even if a model masks padding during attention, the output vectors at padded positions should not automatically be treated as valid sequence content.

Use the attention or validity mask again during pooling. If m_i is 1 for a real token and 0 for padding, masked mean pooling is:

pooled = sum(m_i * h_i) / sum(m_i)

For a batch, the conceptual operation is:

masked = token_vectors * mask[..., None]
summed = sum(masked, token_dimension)
counts = sum(mask, token_dimension)
pooled = summed / counts[..., None]

This is simplified pseudocode, not a framework-specific API. Production code should also define what happens when an example contains zero valid tokens. Dividing by zero is not a meaningful representation; rejecting such input or handling it explicitly is usually clearer than silently producing a value.

Decide which tokens belong in the pool

Padding is not the only token-selection decision. Encoders may insert special tokens for sequence boundaries, classification, separators, or other model-specific purposes.

Whether those tokens should participate in mean pooling depends on the model and training objective. There is no universal rule that special tokens are either meaningful content or harmless noise.

For example, suppose a tokenizer produces:

[SPECIAL] reset password [END]

You could average all four non-padding positions, or average only the two content tokens. Those choices produce different vectors. If the embedding model’s documented pooling procedure includes the special tokens, changing that procedure at inference time can move representations away from the distribution used during training.

For a pretrained embedding model, prefer its documented pooling recipe. Treat a custom pooling rule as a modeling change that should be evaluated, not as a formatting detail.

Special-token pooling relies on training

Some encoder architectures use a designated position as a sequence representation. Pooling then becomes simple selection:

pooled = token_vectors[special_position]

This can be effective when the model was trained so that the selected position receives a useful sequence-level learning signal. The position can attend to other tokens, so its vector may summarize information from across the input.

But the presence of a special token does not by itself guarantee that its final hidden state is a good general-purpose embedding. A model trained for token prediction, sequence classification, or embedding similarity can shape that representation differently.

This leads to a useful rule:

Pooling and training objective should be considered together.

If you train a classifier end to end using one designated token, the classifier can adapt to that representation. If you freeze an arbitrary encoder and use the same token as a semantic embedding without evaluation, you are assuming a property that training may not have encouraged.

Max pooling preserves strong coordinate responses

Max pooling takes the largest value for each embedding dimension across valid tokens:

pooled[j] = max_i h_i[j]

Using the earlier vectors:

h_1 = [1.0, 0.0]
h_2 = [0.5, 1.0]
h_3 = [0.0, 2.0]

max pooling gives:

[1.0, 2.0]

The first component comes from h_1, while the second comes from h_3. The pooled vector therefore need not correspond to any actual token vector.

Max pooling can be useful when strong activations are informative regardless of where they occur. It also has costs: most token values do not directly affect a given pooled dimension, and an unusually large activation can dominate that dimension.

Padding requires special care here too. Multiplying padded vectors by zero is not sufficient when legitimate activations can be negative, because zero could then win the maximum. Invalid positions should instead be excluded from the maximum, commonly by replacing their values with an appropriately low sentinel before reduction.

Learned pooling adds capacity and responsibility

A model can also learn how much each token should contribute. One simple form assigns a score s_i to each valid token, normalizes the scores into weights, and computes a weighted sum:

weight_i = softmax(s_i over valid tokens)
pooled = sum(weight_i * h_i)

This is often called attention pooling. Unlike a fixed mean, it can learn that some contextual token representations matter more for a particular task.

The extra flexibility is not free. Learned pooling introduces parameters, needs task-relevant training data, and can overfit. It also adds computation, although the pooling layer is usually small compared with the encoder itself.

A fixed pooling rule is often preferable when you need a simple baseline, have little training data, or must reproduce a pretrained embedding model’s documented behavior. Learned pooling is more compelling when the downstream task has enough supervision to learn useful token weighting and evaluation shows that the added complexity helps.

Pooling interacts with similarity search

In retrieval systems, the pooled vector is often compared with other pooled vectors. That makes the complete representation pipeline important:

text
 -> tokenize
 -> encode tokens
 -> pool valid token vectors
 -> optional vector normalization
 -> similarity function

Changing only the pooling rule changes every vector in the index. Mixing query vectors produced by one pooling method with document vectors produced by another can make similarity scores difficult to interpret unless that asymmetric design was intentional and trained or evaluated as such.

Pooling also does not solve truncation. If a long document is cut to fit the encoder’s input limit, no pooling operation can recover tokens the encoder never saw. Chunking or another long-document strategy may be necessary before pooling.

For long inputs, mean pooling can also dilute information that appears in only a small part of the sequence. That does not make mean pooling incorrect; it means a single vector may be an inadequate representation for the retrieval granularity you need.

Evaluate pooling as part of the model

Do not choose a pooling method only because it sounds theoretically appropriate. Evaluate the full pipeline on the downstream objective.

For semantic retrieval, measure retrieval quality on representative queries and documents. For classification, compare validation performance while keeping the encoder and training setup controlled. Also measure operational effects if the alternatives change latency or memory use.

A useful comparison might include:

masked mean pooling
model-documented special-token pooling
learned pooling, if training data is available

Keep preprocessing, dataset splits, and evaluation metrics fixed so that the pooling rule is the main changed variable.

If a pretrained model explicitly documents a required pooling and normalization procedure, reproducing that procedure should be the baseline. Benchmarking alternatives is reasonable, but a custom rule creates a different representation system.

Common pooling mistakes

The most common errors come from treating pooling as an interchangeable reduction instead of part of the learned representation pipeline.

Averaging padding tokens. Attention masking and pooling masking solve different stages of the computation. Exclude invalid positions during pooling.

Assuming every special token is a sentence embedding. A position becomes useful for sequence-level tasks because of architecture and training, not because of its token name.

Using zero masking before max pooling. Zero can incorrectly become the maximum when valid values are negative. Exclude invalid positions from the reduction.

Changing pooling without rebuilding stored embeddings. In retrieval, query and indexed document representations must remain compatible with the intended similarity system.

Expecting pooling to fix missing context. Pooling summarizes encoder outputs; it cannot restore truncated text or information the encoder failed to represent.

Choose the simplest pooling rule that matches training

Pooling is the bridge between token-level encoder outputs and sequence-level tasks. Mean pooling gives every valid token equal direct influence, max pooling keeps the strongest coordinate responses, special-token pooling delegates summarization to a designated position, and learned pooling lets training determine token weights.

The practical starting point is straightforward: reproduce the pooling rule expected by a pretrained model when one is specified. For a custom model, begin with a correctly masked simple baseline and compare alternatives on the real downstream metric. The output shape may be identical across methods, but the representation—and therefore the behavior of the system built on it—is not.