How do the linear layers in the attention mechanism work?

Question

I think I now the answer to my question but I dont really get confirmation. When taking a look at the multi-head-attention block as presented in "Attention Is All You Need" we can see that there are three linear layers applied on the key, query and value matrix. And then one layer at the end, which is applied on the output of the matrix multiplication of the score matrix an the value.

The three linear layers at the beginnging: When the key/query/value with shape (seq-len x emb-dim) enter the linear layer the output is still (seq-len x emb-dim). Does that mean, the same linear layer is applied on every "index" of the input matrix. Like this (pseudo-code):

fc = linear(emb-dim, emb-dim) # in-features and out-features have the shape of emb-dim
output_matrix = []

for x in key/query/value:
    # x is one row the input matrix with shape (emb-dim x 1)
    x = fc(x)
    # x after the linear layer has still the shape of (emb-dim x 1)
    output_matrix.append(x)

# output_matrix has now the shape (seq-len x emb-dim); the same as the input-matrix

So is this indeed what happens? I couldn't explain why the output is the same as the input otherwise.

The linear layer before the output: So the output of the matrix multiplication of the score matrix an the value is also (seq-len x emb-dim) and therefore the output of the linear layer is too. So the output of the whole attention block has the same shape as the input.

So Im just asking for comfirmation if the explaination I wrote is correct. And if not: What am I understanding wrong?

Extra question: When I want to further use the output of the attention block further for classification, I would have to take the mean along the seq axis in order to get a vector of fixed shape (emb-dim x 1) so I can feed it into a classification layer. But I guess that valueable information is getting lost in that process. My question: Could I replace the last linear layer with an RNN to get the desired output shape and without losing information?

Answer 1

Your understanding is not correct.

The relevant information is described in the original paper in section 3.2.2:

The three sets of projection matrices you are referring to are $W^Q_i \in \mathbb{R}^{d_{model} \times d_k}$ for the Queries, $W^K_i \in \mathbb{R}^{d_{model}\times d_k}$ for the Keys and $W^V_i \in \mathbb{R}^{d_{model}\times d_v}$ for the Values. Notice that the $i$ subindex in the matrix names refers to the attention head, indicating that there is a different matrix for each attention head. The final projection matrix is $W^O \in \mathbb{R}^{hd_v \times d_{model}}$.

Given the number of attention heads, $h = 8$, the dimensions of the matrices are defined by $d_k=d_v=d_{model}/h=64$.

The three sets of matrices project the embedding dimensionality $d_{model}$ into a space with 8 times smaller dimensionality ($d_k=d_v=d_{model}/8$). However, note that for each of $W^K$, $W^V$ and $W^Q$ there are 8 matrices (one per attention head) and, analogously, 8 scaled dot products are computed. The results of the dot products are 8 vectors of dimensionality $d_{model}/8$; those 8 vectors are concatenated (see figure below), giving a tensor with the original dimensionality $d_{model}$.

Then, the final matrix multiplication by $W^O$ doesn't change the dimensionality, obtaining again the original one.

---

About your second question, either of the approaches you describe (averaging and using an RNN) is technically feasible, but what people normally do when using transformers for classification is to use BERT's approach, that is, adding a special token [CLS] at the beginning of the sequence and using the output at that position for the classification, just projecting with a matrix multiplication into a space with dimensionality equal to the number of classes.