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?

enter image description here

With Eigen Tensor I have these incorrect duplicate results

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

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

Eigen::Tensor<double, 3> key(2, 2, 3);
                         {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]]]]

1 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.
  • $\begingroup$ I was hoping there'd be an Eigen equivalent of numpy and TF's einsum, specifically something like "abc,adc->abd". They have a merge request open for it (gitlab.com/libeigen/eigen/-/merge_requests/52/…) but it's a big feature and, according to discussion, they're putting it off util after version 3.4. $\endgroup$
    – rkuang25
    May 18, 2023 at 0:44
  • $\begingroup$ I cant help you with imementations in Eigen unfortunately. I can pnly describe the theoretical answer to your question, which i hope I have done. For implementation questions, the DS SE is a better place to start. $\endgroup$ May 18, 2023 at 6:02
  • $\begingroup$ The theoretical portion is simple, it's just the implementation in Eigen that I'm having difficulty. I have opted for a for-loop with omp parallel instead. For large enough tensors, I found that a for-loop with omp parallel and chip/slice is also nearly as fast. $\endgroup$
    – rkuang25
    May 18, 2023 at 6:05
  • $\begingroup$ The question "How Does Self-Attention Query and Key Multiplication Work For Batched Inputs?" and i get the idea that that is answered. If you'd want advice on implementation the DSSE is a better place to open up a question. $\endgroup$ May 18, 2023 at 6:54
  • $\begingroup$ Thank you Robin. I have arrived at a solution already. $\endgroup$
    – rkuang25
    May 18, 2023 at 7:12

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .