Sentence Embedding (Masked Mean-Pool)

Background

A bi-encoder's sentence embedding starts with per-token vectors. Mean pooling with attention mask is the dominant aggregation: average only the real-token positions (not padding), then L2-normalise so the embedding lives on the unit sphere for cosine-similarity retrieval.

Problem statement

Implement masked_mean_pool(token_embs, attention_mask):

$$\mathbf{s} \;=\; \frac{\sum_t m_t\,\mathbf{e}_t}{\sum_t m_t},\qquad \hat{\mathbf{s}} \;=\; \frac{\mathbf{s}}{\|\mathbf{s}\|_2}$$

with $m_t \in \{0, 1\}$. Return the L2-normalised sentence embedding.

Input

• token_embs - shape $(T, D)$, per-token vectors.
• attention_mask - shape $(T,)$ of $\{0, 1\}$.

Output

• np.ndarray of shape $(D,)$, unit norm.

Examples

Input: token_embs = [[1,0],[0,1],[0,0]], mask=[1,1,0]
Output: mean=[0.5, 0.5]; L2 norm -> [1/sqrt(2), 1/sqrt(2)]

Constraints

• Mask sum must be > 0; otherwise ValueError.
• If the mean vector is all-zero, return a zero vector rather than dividing by 0.

Notes

• vs CLS-pool. Encoder-style BERT uses the CLS token; sentence-transformers / E5 use mean-pool. Mean-pool is the de-facto choice for embedding retrieval.