Transformer model of the original Attention paper has a decoder unit that works differently during Inference than Tranining.

I'm trying to understand the shapes used during decoder (both self-attention and enc-dec-attention blocks), but it's very confusing. I'm referring to this link and also the original Attention paper

In Inference, it uses all previous tokens generated until that time step (say kth time-step), as shown in the diagram below and explained at this link.

Another diagram that shows self-attention and enc-dec-attention within decoder:

enter image description here


However when I look at actual shapes of the QKV projection in the decoder self-attention, and feeding of the decoder self-attention output to the "enc-dec-attention"'s Q matrix, I see only 1 token from the output being used.

Let's assume 6 deocder blocks one after the other in the decoder stack (which is the base transformer model).

I'm very confused how the shapes for all matrices in the Decoder blocks after decoder-1 of the decoder-stack (more specifically decoder-block-2 decoder-3, decoder-4..., decoder-6 of the decoder stack) self-attention and enc-dec-attention can match up with variable length of input to the decoder during inference. I looked at several online material but couldn't find answer. I see only the BGemms in the decoder's self-attention (not enc-dec-attention) using the variable shapes until all previous k steps, but all other Gemms are fixed size.

  • How is that possible? Is only 1 token (last one from decoder output) is being used for qkv matmuls in self-attention and Q-matmul in enc-dec-attention (which is what I see when running the model)?
  • Could someone elaborate how all these shapes for QKV in self-attention and Q in enc-dec-attention match up with decoder input length being different at each time-step?**
  • $\begingroup$ Hello. Rather than saying that you're confused in the title, could you please put your specific question there? Thanks! $\endgroup$
    – nbro
    Nov 9, 2021 at 16:13
  • 1
    $\begingroup$ Updated to specific question in the title per your comment. $\endgroup$
    – Joe Black
    Nov 10, 2021 at 2:59

1 Answer 1


Edit 3

OP seems to think value, query and keys are supposed to be different in the original Vaswani multi-head attention. As can be seen in Keras' documentation on their implementation of the multi-headed attention layer, "If query, key, value are the same, then this is self-attention."

Edit 2

One thing missing from the graphics you use are the skip connections in transformers. Look at figure 1 in the original Vaswani et al paper. The skip connections should make it quite obvious what the shape of the outputs have to be like after each layer; After all, you cannot add two tensors that do not have the same shape.


I realize now that your question is regarding the key, value and query values in an attention mechanism. They are always the same. It's called self-attention for that reason. The attention mechanism used in all papers I have seen use self-attention: K=V=Q

Also, consider the linear algebra involved in the mechanism; The inputs make up a matrix, and attention uses matrix multiplications afterwards. That should tell you everything regarding the shape those values need.

Here's a valuable visual that depicts attention in a slightly different way: self-attention

It's from this blog post which explains self-attention in-depth.

Transformers only output one prediction at a time because it is an autoregressive model.

Lets break a transformer during runtime step by step. The decoder gets the output as its input, which is done in the following manner:

Step one: The decoders input is only the start token and padding $(start,0,0,..,0,0)$

Step two: The decoder input is the start-token and the prediction from step one + padding $(start,y_1,0,0,..,0,0)$

Step three: The decoder input is the decoders input in step two and the prediction from step two + padding $(start,y_1,y_2,0,0,..,0,0)$

And, so on..

So, if you want n outputs, the transformer will have to run n times. If you want to train it to stop by itself after a while, you can introduce an end-of-sequence token which stops inference once it's predicted by the transformer. The sequences that are being fed into the decoders and encoders are always max_length long. That length is preserved using padding.

  • $\begingroup$ I understand how the decoder does prediction in time-steps start, y1, y2...etc unlike during Training (teacher-forcing), but my question was more about inputs to the K/Q/V of EACH decoder block in the "decoder-stack". I updated the title to more specific question, per another comment. Assume 6 decoder blocks in the decoder-stack. So the input to the 1st decoder (decoder-1) block's Q/K/V will be (start, y1, y2,... yn) at time-step n. But what'll be the input to Q/K/V matmuls of self-attention of decoder-2 in the decoder-stack at time n? $\endgroup$
    – Joe Black
    Nov 10, 2021 at 2:57
  • $\begingroup$ I also understand the K/V of the "enc-dec-attention" (from 2nd diagram in OP) of a given decoder block (say decoder-2) of the decoder-stack would come from encoder outputs. But my question is specifically about the self-attention inputs in decoder blocks of the decoder stack? what would be their shapes/inputs? $\endgroup$
    – Joe Black
    Nov 10, 2021 at 2:59
  • $\begingroup$ Added more details about 6 decoder-blocks in the decoder stack in the OP body. Also, for decoder-1 would all 3 Q/K/V inputs be (start, y1, y2, ....yn) or would Q only consist of the last token yn but K/V would consist of all tokens till n i.e. (start, y1, y2, ....yn) ? $\endgroup$
    – Joe Black
    Nov 10, 2021 at 3:05
  • $\begingroup$ I think I understand that part already, except I'm not sure about K=Q=V for decoder. What's the input to decoder-2 (in decoder-stack after decoder-1) K, Q, V? Since the output of decoder-1 is a single-token, assuming hidden-size of 1024 and batch-size of 1, it'd be 1 tensor of shape 1x1024 being output. Are you saying 1x1024 is being fed to K,Q, V of decoder-2 whereas nx1024 was fed to K,Q,V of decoder-1? $\endgroup$
    – Joe Black
    Nov 10, 2021 at 17:27
  • $\begingroup$ Please note my question specifically is about input to EACH decoder in decoder-stack (1 to 6 decoder blocks), and notice how that can cause issues. $\endgroup$
    – Joe Black
    Nov 10, 2021 at 17:27

You must log in to answer this question.

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