Embedding Models: Architecture, Pooling & Selection › Worked Example: Mean Pooling Step by Step

To make mean pooling concrete, let's walk through a numerical example with a toy model. Suppose we have a toy transformer with hidden dimension $d = 4$, and we encode the three-token sequence ["The", "cat", "meows"] with one padding token added to match a batch length of 4.

After the final transformer layer, suppose the hidden states are as follows (each row is one token's hidden state, each column is one of the four dimensions):

$$\mathbf{H} = \begin{pmatrix} 0.8 & -0.2 & 0.5 & 0.1 \\ -0.3 & 0.9 & 0.4 & -0.6 \\ 0.6 & 0.1 & -0.7 & 0.3 \\ 0.0 & 0.0 & 0.0 & 0.0 \end{pmatrix}$$

where the first row is the hidden state for "The", the second for "cat", the third for "meows", and the fourth for the padding token. The attention mask is $\mathbf{m} = [1, 1, 1, 0]$, showing that the first three tokens are real and the fourth is padding.

Step 1: Apply the mask. Multiply each hidden state row by its mask value:

$$m_1 \cdot \mathbf{h}_1 = 1 \cdot (0.8, -0.2, 0.5, 0.1) = (0.8, -0.2, 0.5, 0.1)$$ $$m_2 \cdot \mathbf{h}_2 = 1 \cdot (-0.3, 0.9, 0.4, -0.6) = (-0.3, 0.9, 0.4, -0.6)$$ $$m_3 \cdot \mathbf{h}_3 = 1 \cdot (0.6, 0.1, -0.7, 0.3) = (0.6, 0.1, -0.7, 0.3)$$ $$m_4 \cdot \mathbf{h}_4 = 0 \cdot (0.0, 0.0, 0.0, 0.0) = (0.0, 0.0, 0.0, 0.0)$$

Step 2: Sum the masked hidden states across all token positions, dimension by dimension:

$$\sum_{i=1}^{4} m_i \cdot \mathbf{h}_i = (0.8 + (-0.3) + 0.6 + 0.0,\; -0.2 + 0.9 + 0.1 + 0.0,\; 0.5 + 0.4 + (-0.7) + 0.0,\; 0.1 + (-0.6) + 0.3 + 0.0)$$ $$= (1.1,\; 0.8,\; 0.2,\; -0.2)$$

Step 3: Divide by the number of real tokens. The sum of the mask values is $\sum_{i=1}^{4} m_i = 1 + 1 + 1 + 0 = 3$, so we divide each dimension by 3:

$$\mathbf{e} = \frac{1}{3}(1.1,\; 0.8,\; 0.2,\; -0.2) = (0.367,\; 0.267,\; 0.067,\; -0.067)$$

Step 4: L2-normalize. Compute the Euclidean norm of this vector:

$$\|\mathbf{e}\|_2 = \sqrt{0.367^2 + 0.267^2 + 0.067^2 + (-0.067)^2}$$ $$\begin{aligned} &= \sqrt{0.1347 + 0.0713 + 0.0045 + 0.0045} \\ &= \sqrt{0.2150} \\ &\approx 0.464 \end{aligned}$$

Then normalize:

$$\hat{\mathbf{e}} = \frac{\mathbf{e}}{\|\mathbf{e}\|_2} = \frac{(0.367, 0.267, 0.067, -0.067)}{0.464} \approx (0.791, 0.575, 0.144, -0.144)$$

The key insight from this example is that mean pooling is a simple weighted average where the weights come from the attention mask, and the normalization step ensures that the output vector lies on the unit sphere. The entire computation is differentiable with respect to the hidden states $\mathbf{H}$, which means gradients can flow back through the pooling operation during training, teaching the transformer to produce hidden states that aggregate well when averaged.