LLM Systems
Exclude padding when mean-pooling token embeddings
Pool included token vectors and divide by their actual count. A synthetic NumPy fixture catches padding-dependent means and rejects empty masks.
Evidence and testing
Review record: Tensor Almanac automated evidence gate (no independent human review). Example versions, recorded outputs and limitations appear below.
A batch mean includes every padded position
When converting token vectors into one vector per sequence, averaging the complete token axis makes padding part of the calculation. The all-MiniLM-L6-v2 model card gives an explicit mask-aware mean-pooling recipe: include the selected token vectors in the numerator and use the included token count in the denominator. Passing an attention mask to the encoder does not replace that explicit pooling step in its recipe. [1]
The executed fixture contains two synthetic sequences padded to four token positions. Their included-token counts are two and three. Large nonzero padding vectors make the error obvious: ordinary means are (50.5, 50.5) and (-6, -6), while the intended means are (1, 1) and (2, 2). The numbers are deliberately artificial and are not transformer outputs.
References: [1]
Reduce the token axis and retain one count per sequence
The helper accepts arrays shaped (batch, tokens, features) and a binary mask shaped (batch, tokens). It expands the mask along the feature dimension, replaces excluded vectors with zero, sums along the token axis and divides by one included-token count per sequence. NumPy documents axis reduction and keepdims; retaining the count dimension lets each feature of a sequence use that sequence's denominator. [2]
Run python examples/masked_token_pooling.py. Assertions compare the pooled vectors with hand-calculated results and verify that appending two additional masked positions leaves both vectors unchanged. A second check replaces padding vectors with zero. The ordinary mean is still wrong because it divides by four positions instead of the two or three included tokens.
The example checks a single included token and rejects six invalid fixtures: an entirely excluded batch, a mask-shape mismatch, a nonbinary mask, incorrect embedding rank, an empty token axis and non-finite embedding values. It rejects any sequence with no included tokens rather than inventing a sentence representation. That is an explicit helper policy, different from the model card's denominator clamp.
Executed with Python 3.12.14 and NumPy 2.3.5 on synthetic CPU arrays. No tokenizer or model was loaded. Exact stdout and source-byte SHA-256 are preserved in private evidence.
"""Audit padding exclusion in mean pooling of synthetic token vectors."""
import numpy as np
def masked_mean(hidden, mask):
hidden = np.asarray(hidden, dtype=np.float64)
mask = np.asarray(mask)
if hidden.ndim != 3 or any(size == 0 for size in hidden.shape):
raise ValueError("nonempty batch-token-feature array required")
if mask.shape != hidden.shape[:2] or not np.all(np.isin(mask, [0, 1])):
raise ValueError("aligned binary token mask required")
if not np.all(np.isfinite(hidden)):
raise ValueError("finite token vectors required")
counts = np.sum(mask, axis=1, keepdims=True, dtype=np.float64)
if np.any(counts == 0):
raise ValueError("each sequence must contain an included token")
included = np.where(mask[:, :, None].astype(bool), hidden, 0.0)
return np.sum(included, axis=1, dtype=np.float64) / counts
hidden = np.array([[[2, 0], [0, 2], [100, 100], [100, 100]],
[[3, 0], [0, 3], [3, 3], [-30, -30]]], dtype=np.float64)
mask = np.array([[1, 1, 0, 0], [1, 1, 1, 0]])
pooled = masked_mean(hidden, mask)
naive = np.mean(hidden, axis=1)
assert np.array_equal(pooled, [[1, 1], [2, 2]])
assert np.array_equal(naive, [[50.5, 50.5], [-6, -6]])
extra = np.concatenate([hidden, np.full((2, 2, 2), 999.0)], axis=1)
extra_mask = np.concatenate([mask, np.zeros((2, 2), dtype=int)], axis=1)
assert np.array_equal(masked_mean(extra, extra_mask), pooled)
zero_padded = np.where(mask[:, :, None].astype(bool), hidden, 0.0)
assert np.array_equal(masked_mean(zero_padded, mask), pooled)
assert not np.allclose(np.mean(zero_padded, axis=1), pooled)
assert np.array_equal(masked_mean(hidden[:1, :1], [[1]]), [[2, 0]])
rejected = 0
for h, m in ((hidden, np.zeros_like(mask)), (hidden, [[1, 0], [1, 1]]),
(hidden, [[1, 2, 0, 0], [1, 1, 1, 0]]),
(hidden[:, 0], mask), (hidden[:, :0], mask[:, :0]),
(np.full_like(hidden, np.nan), mask)):
try:
masked_mean(h, m)
except ValueError:
rejected += 1
assert rejected == 6
print("unmasked means:", naive.tolist())
print("included token counts:", mask.sum(axis=1).tolist())
print("masked means:", pooled.tolist())
print("adding masked padding preserves result:", np.array_equal(masked_mean(extra, extra_mask), pooled))
print("zero padding still biases an unmasked mean:", not np.allclose(np.mean(zero_padded, axis=1), pooled))
print("invalid pooling fixtures rejected:", rejected)
Recorded output
unmasked means: [[50.5, 50.5], [-6.0, -6.0]] included token counts: [2, 3] masked means: [[1.0, 1.0], [2.0, 2.0]] adding masked padding preserves result: True zero padding still biases an unmasked mean: True invalid pooling fixtures rejected: 6
Match the pooling contract of the embedding model
The supplied mask defines which token vectors participate. Decisions about special tokens, truncation and normalization must match the intended embedding model; a generic mask-aware mean is not a universal sentence-embedding recipe. This fixture isolates padding exclusion and does not reproduce the full MiniLM encoding pipeline.
The helper handles small finite float64 arrays and binary masks. It allocates a masked intermediate, does not implement weighted pooling and does not measure retrieval quality or throughput. Use padding invariance as a targeted pooling test while checking the actual encoder and tokenizer separately.