Dimension Mismatch in Transformer Decoder: What Are the Input and Output Dimensions?

Question

My understanding is that in the decoder, the output of the masked self-attention mechanism is expected to have dimensions (o_len,d_model), where o_len is the current output length.

However, an issue arises when the keys (K) and values (V) used in the self-attention of the decoder are obtained from the output of the encoder. Their dimensions are (n, d_model), where n represents the number of embedding vectors. This poses a problem during the computation of Q x K^T because Q has size (o_len,dq), while K^T has a size of(d_model,n), and dq is not equal to dmodel. ​

Answer 1

Q and K must have the same embedding dimension in order to compute scalar products, meaning dq must equal dmodel in your case.

However, the PyTorch MultiheadAttention class accepts different embedding dimensions for queries (embed_dim) and keys (kdim) because PyTorch implements the projection matrix internally.

see the original Attention Is All You Need paper