How Does The Scaled Dot Product's Dimensions Work Out In Mult-Head Attention?

Question

I don't understand how self-attention works with batched values for the $Q \times K^T $ step. According to the diagram below (assume 1 head), once we get past the first 3 linear steps, we arrive at the equation. $$softmax(\frac{Q \, K^T}{\sqrt{d_k}})V$$ Assuming Q and K have dimensions $[\text{batch}, \text{sequence}, \text{key_dim}]$, if the tensor contraction goes like $$Q \, K^T = [\text{batch}_q, \text{sequence}_q, \text{query_dim}_q] \times [\text{batch}_k, \text{key_dim}_k, \text{sequence}_k] = [\text{batch}_q, \text{batch}_k, \text{query_dim}_q, \text{key_dim}_k]$$ there is an extra batch dimension in the result.

I can avoid this extra batch dimension by doing a for-loop over the batch dimension and compute the formula 1-to-1 for each query and key. But one of the benefits of self-attention is that it's parallelizable so there must be a way to do this without a for-loop. How do I achieve this?

With Eigen Tensor I have these incorrect duplicate results

Eigen::Tensor<double, 3> query(2, 2, 3);
query.setValues({{
                         {1, 6, 9},
                         {8, 7, 11}},

                 {       {1, 6, 9},
                         {8, 7, 11}}});

Eigen::Tensor<double, 3> key(2, 2, 3);
key.setValues({{
                         {4, 5, 0},
                         {10, 15, 7}},

                 {       {4, 5, 0},
                         {10, 15, 7}}});

// contract along last dim of query and key
Eigen::array<Eigen::IndexPair<int>, 1> matmul{Eigen::IndexPair<int>(2, 2)};

Eigen::Tensor<double, 4> qk_t = query.contract(key, matmul);

std::cout << "QxK_T: dimensions" << qk_t.dimensions() << "\n" << qk_t << "\n";

QxK_T: [2, 2, 2, 2]
[[[[ 34, 163],
   [ 34, 163]],

  [[ 67, 262],
   [ 67, 262]]],


 [[[ 34, 163],
   [ 34, 163]],

  [[ 67, 262],
   [ 67, 262]]]]

Answer 1

Say you have 3-dimensional tensors Q & K of shape (B, T, C) (for Batch, Time, and Channels).

1. K is transposed. 'Transposing' a 3-d tensor in this case means switching around the last two dimensions. Hence shape(K.T) = (B, C, T).
2. Q and K.T are matrix multiplied. Matrix multiplication is done on a per-sample basis. Logically, samples are not mixed. Hence we get matmul(Q, K.T) gives the matrix multiplication (B, T, C) matmul (B, C, T) -> (B, T, T)
   • Think of it like this. Take the batch apart into individual samples, and simply matrix-multiply the individual samples.