Consider the decoder part of the popular Transformer architecture; briefly put, the decoder module consists of a composition of self-attention layers and performs auto-regressive prediction. Because of the masked attention, the output at any time step $j$ depends only on the previous tokens, from $1$ to $j$. Mathematically, a simplified one self-attention layer would look like:

$$ \mathbf{y}_j = \mathrm{softmax}\left(\sum_{i=1}^j \left\langle \mathbf{k}(\mathbf{x}_i), \mathbf{q}(\mathbf{x}_j) \right\rangle \mathbf{v}(\mathbf{x}_i)\right) $$

However, in order for the model not to cheat and just copy the input to the output, the input sequence is shifted to the right with one token. So, the output $\mathbf{y}_j$ corresponds in fact to the prediction of the next token, $\hat{\mathbf{x}}_{j+1}$:

$$ \hat{\mathbf{x}}_{j+1} = \mathrm{softmax}\left(\sum_{i=1}^j \left\langle \mathbf{k}(\mathbf{x}_i), \mathbf{q}(\mathbf{x}_j) \right\rangle \mathbf{v}(\mathbf{x}_i)\right) $$

If we consider a concrete example

enter image description here

the probability for the word "sat" given the previous words is

$$ p(\mathrm{sat} | \mathsf{s}, \mathrm{the}, \mathrm{cat}) = \mathrm{softmax}\left(\alpha(\mathsf{s}, \mathrm{cat}) \mathbf{v}(\mathsf{s}) + \alpha(\mathrm{the}, \mathrm{cat}) \mathbf{v}(\mathrm{the}) + \alpha(\mathrm{cat},\mathrm{cat})\mathbf{v}(\mathrm{cat})\right)_{\mathrm{sat}} $$

where $\mathsf{s}$ and $\mathsf{e}$ denote the start and end tokens, and $\alpha$ the key-query inner product.

Given the above, it seems to me that the Transformer decoder pools the features based on the previous token, the $j$-th one, to predict the current one, the $j +1$-th.

  1. Is my understanding correct?
  2. What is the motivation for using the previous token as a way to aggregate the features for the current token? I could think of alternatives to using the previous token $\mathbf{x}_j$ to build the query: use the positional encoding $\mathbf{p}_{j+1}$ or average all the previous vectors $\mathbf{x}_1, \dots, \mathbf{x}_j$.

1 Answer 1


Your interpretation would be correct if the decoder would have a single self-attentive layer with a single head.

In the first layer, you indeed use the previous token query the history and the encoder, but you cannot know how much from the word is actually used: the linear projection for queries, keys, and values, can keep only a little information and do some more complex pooling over the states. After the first self-attentive layer, the context vector can already contain information from the entire sentences, and the query word can play a minor role, more important can be the content of the previous states, it depends on the projection.

In the following layers, the query vector is indeed something hardly interpretable, but it is a complex combination of all encoder states, previous decoder states which already went through non-linear layers.

I think both the things you suggest might work, but they probably would not make much difference. I think the reason pro using the previous token is not that it is an extremely good query, but that you gradually create the target sentence on the input to be used as keys and values. To keep it simple the last token can also be the query.

  • $\begingroup$ Thank you very much for the answer! You are right—if we assume multiple layers then at the second layer the new query will be a (complex) mix of the previous tokens (in essence, similar to the second idea that I mentioned). And I also tend to agree with your intuition that the alternatives I’m suggesting won’t probably make a big difference in terms of final results. $\endgroup$ Commented Apr 6, 2020 at 9:24
  • $\begingroup$ You are saying that the interpretation is correct if the decoder has “a single self-attentive layer with a single head”. Wouldn’t the statement be also valid if had a single self-attentive layer, but multiple heads? $\endgroup$ Commented Apr 6, 2020 at 9:25

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.