Masked Mean Pooling

Masked mean pooling turns many token embeddings into one embedding that represents an entire sentence or document. For example:

• classifying an entire document
• creating one embedding for semantic search
• finding the cosine similarity between two sentences or documents

A model often produces one vector per token, $x: [B, T, D]$. Pause: what do these dimensions represent?

Mean pooling converts this into a tensor of shape $[B, D]$, and padding must not affect the result.

Let $m: [B, T]$ be our mask, which tells us which positions are real and which are padding.

The numerator

• $m_{b,t}$ equals 0 if the position is padding, 1 if it's a real token.
• $x_{b,t}$ is a vector of length $D$ — the token embedding.

The denominator

• $\max(\text{sum}, 1)$ just avoids dividing by 0.
• $\sum_t m_{b,t}$ gives the number of valid tokens, so dividing gives us the "mean".

In code:

def masked_mean_pool(x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """Mean-pool valid tokens from x [B, T, D] using mask [B, T]."""
    # x    : [B, T, D]
    # mask : [B, T]
    B, T, D = x.shape
    mask_unsqueezed = mask.unsqueeze(-1)      # [B, T] -> [B, T, 1]
    masked_x = x * mask_unsqueezed            # [B, T, D] * [B, T, 1] -> [B, T, D]
    numerator = masked_x.sum(dim=1)           # [B, T, D] -> [B, D]
    mask_sum = mask.sum(dim=1).clamp(min=1)   # [B, T] -> [B]
    denominator = mask_sum.unsqueeze(-1)      # [B] -> [B, 1]
    return numerator / denominator            # [B, D] / [B, 1] -> [B, D]