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