2
$\begingroup$

Forgive me the long introduction, but I want to be very specific with my question.

The Annotated Transformer implements the Transformer architecture from "Attention Is All You Need". The MultiHeadedAttention class implements the multi-headed attention mechanism with what, to me, look like tricks to simplify the code and improve performance. Specifically, the parallel application of multiple attention heads is implemented by handling all heads in the same matrix, one for each step of the attention computation.

My question is related to this line:

query, key, value = \
            [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
             for l, x in zip(self.linears, (query, key, value))]

It consists of:

  1. applying a linear layer l to the input x
  2. transforming the output to a 4D Tensor with dimensions describing, respectively: batch, attention head, word in the phrase, projection dimension

The question is: why does it have to be a combination of view() and transpose() rather than just a single call to view()? I could implement that line as follows:

query, key, value = \
            [l(x).view(nbatches, self.h, -1, self.d_k)
             for l, x in zip(self.linears, (query, key, value))]

It produces the output of the same shape and I checked that it works (at least on a CPU; I don't know about CUDA).

So why is it done this way? What is its advantage?

$\endgroup$

2 Answers 2

1
$\begingroup$

The reason is technical: view can only operate on contiguous tensors and is supposed to be faster, transpose works for both contiguous and non-contiguous tensors.

There is a blog post that tries to explain that, for a better understanding of what contiguous tensor means, there is a nice PyTorch forum answer.

$\endgroup$
1
  • $\begingroup$ While these are good facts to know about view and transpose, they do not address OP's question, which is about how to compute multi-head attention in a transformer. Specifically, applying view(nbatches, self.h, -1, self.d_k) splits tensors along the sequence dimension, which is incorrect. The desired result is to split on the embedding dimension. $\endgroup$
    – Sycorax
    Commented Sep 17, 2023 at 15:50
1
$\begingroup$

The reason is not technical. You apply the transpose because you need to transpose the axes of the tensors.

Your input $x$ has shape $B \times T \times D$. Once you forward through the linear layer l you get the query embeddings which have the same shape. What you want to do is split this tensor along the last dimension, because this is where the embedding dimension is. That is why you have to do l(x).view(B, T, n_h, D//n_h). After that you transpose axes 1 and 2 to get the desired shape.

If you apply l(x).view(B, n_h, T, D//n_h) you split the tensor along the wrong dimension - the sequence length dimension. Components of the embeddings of different tokens will get mixed up. You get the correct shape, but the wrong numbers.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.