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

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]]]]


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)