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