Why does a decoder generate all hidden states during inference?

Question

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]

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.