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
    Commented Nov 9, 2021 at 16:13
  • 1
    $\begingroup$ Updated to specific question in the title per your comment. $\endgroup$
    – Joe Black
    Commented Nov 10, 2021 at 2:59
  • $\begingroup$ Please see this (newer) post for an answer to your question $\endgroup$ Commented May 10, 2023 at 9:52


You must log in to answer this question.