Attention step 3 — weighted sum of values

Background

The final attention step turns the weights into an output by taking a weighted sum of the value vectors:

$$\text{output} = \alpha V, \qquad \text{output}_i = \sum_j \alpha_{ij}\,\mathbf{v}_j$$

α is (n_q, n_k) (rows sum to 1), V is (n_k, d_v), and the output is (n_q, d_v) — each query's answer is a convex combination of the values.

Problem statement

Implement aggregate(weights, V) returning $\alpha V$.

Input

• weights — array-like of shape (n_q, n_k), attention weights.
• V — array-like of shape (n_k, d_v), value vectors.

Output

Returns an np.ndarray of shape (n_q, d_v).

Examples

Example 1 — the lesson's worked weights

Input:  weights = [[0.506, 0.186, 0.308]],
        V = [[10, 0], [0, 5], [2, 2]]
Output: [[5.676, 1.546]]

Explanation: mostly $v_1$, a touch of $v_3$, almost none of $v_2$.

Constraints

• Compute weights @ V.
• Return shape (n_q, d_v).

Notes

• Because the weights are a probability distribution, the output is a blend of values — never a hard pick — which keeps the whole operation differentiable.