Seems that in Vanilla transformers at least (a la AIAYN), during inference time, the hidden states are generated for all tokens in the input sequence, but only the last one is used to predict the next token.

My question - why are the other hidden states produced at all during inference? Wouldn't it be more efficient if only the last hidden state was produced?

For example, here : https://nlp.seas.harvard.edu/annotated-transformer/#greedy-decoding

out = model.decode(
            memory, src_mask, ys, subsequent_mask(ys.size(1)).type_as(src.data)
        prob = model.generator(out[:, -1])
        _, next_word = torch.max(prob, dim=1)
        next_word = next_word.data[0]

1 Answer 1


I understand the you mean that only the hidden states at position t - 1 of the decoder are used to predict the token at position t. This is not true. To compute the hidden states at position t of layer L, the hidden state/token from positions 1 to t - 1 of layer L - 1 are used (all previous positions from the previous layer are used).

The example you posted is a part of a greedy decoding implementation. This should ho in a loop to autoregressively decode token by token. It also misses the a line from the source that updates ys concatenating it with the predicted token so that, at the following decoding step, the model uses all previously decoded tokens.

Note that in this basic implementation you are recomputing the hidden states of the first positions at every loop iteration. This wasted computtion can be saved by caching these hidden states and giving them to the decoding for it to reuse them.

  • $\begingroup$ Thanks. I understand it uses the previous ys (for attention calculation) and auto-regressively updates. I just don't see where it uses any of the hidden states of the previous ys. They seem to all be produced in parallel and then only the last is used. Can you show me where they are used please? $\endgroup$
    – dashnick
    Apr 8, 2023 at 10:38
  • $\begingroup$ In the forward method of the Decoder class you can see loop that compues each layer. These are the hidden states of all positions up to the current position. $\endgroup$
    – noe
    Apr 8, 2023 at 10:54
  • $\begingroup$ Thanks again. I see they are computed, but I don't see how they inform out[:,-1] / the last hidden state. Since only the last hidden state is used, any calculation that doesn't affect this last hidden state is a waste. I know I'm missing something :) $\endgroup$
    – dashnick
    Apr 8, 2023 at 11:18
  • $\begingroup$ I think it's for the attention over the encoder encodings, but don't know why these are necessary then in decoder only models $\endgroup$
    – dashnick
    Apr 8, 2023 at 11:25
  • $\begingroup$ The whole tensor with all the positions is handled together, so you won’t find any loop over the positions. The result for all positions is computed at once in parallel. The encoder-decoder attention is structurally the same as the decoder-only attention; the difference is that the query comes from the encoder last layer output and the keys and vales from the current decoder layer, while in the decoder-only both keys, values and queries come from the current decoder layer. $\endgroup$
    – noe
    Apr 8, 2023 at 12:27

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.