Does the transformer decoder reuse previous tokens' intermediate states like GPT2?

Question

I recently read Jay Alammar's blogpost about GPT-2 (http://jalammar.github.io/illustrated-gpt2/) which I found quite clear appart from one point : He explains that the decoder of GPT-2 processes input tokens one at a time, only actively processing the last input token, the past tokens being saved in memory already and "passively" reused without reevaluation.

From my understanding of the transformer architecture, I had the impression that the decoder reevaluates every token generated at each generation. Is this then a difference between the decoder from GPT-2 or does the decoder from a "classical' transformer also work this way ?

Intuitively I would think that it would make more sense the reevaluate everything at each iteration since new dependencies between words can appear that weren't there at the beginning and would then not be taken into account if past processed words are passively reused.

I hope I am making sense, can someone with knowledge about the GPT2 architecture help me clarify this ?

Answer 1

My understanding is that transformer decoders and transformer encoder-decoder models typically operate in the way that the GPT-2 does, i.e., representations in the generated sequence are computed once and then reused for future steps. But you are correct that this is not the only way things can be done. One could recompute the representations for all tokens in the partially-generated sequence using full self-attention over the tokens in the sequence generated so far (there's no mathematical hindrance to doing this -- it's akin to running a typical transformer encoder over the sequence of words in the partially-generated sequence).

But this additional computation is not commonly done as far as I can tell from the literature. I think there are at least two reasons. First, as noted by others, it's cheaper computationally to use previously-computed representations from earlier time steps (though it leads to different results, and I have not seen an empirical comparison in any papers). Second, it matches how training is done. During training, a consequence of masking in self-attention is that the representation at output position i is computed using representations at output positions <= i. That means that during training, there is only a single representation computed for output position i for each layer. That matches what happens at inference time using the standard approach that we've been discussing and which is used in the GPT-2.

If we wanted to train a model in which the representation for an output position was computed based on all available output representations (always excluding the ones that have not yet been "generated" of course), then we would need to compute multiple representations for each output position during training, one for each possible uncovered partial right context. For example, if we train a language model on windows of size 512, we would have to compute (about) 512 representations for the first word, one corresponding to the loss for generating each subsequent word in the window. This would lead to a very large computation graph and slow down training. However, it may work better because it leads to richer output representations, so please try it and let us know. :)

Answer 2

The past token internal states are reused both in GPT-2 and any other Transformer decoder.

For example, in fairseq's implementation of the transformer, these previous states are received in TransformerDecoder.forward in parameter incremental_state(see the source code).

Remember that there is a mask in the self-attention blocks in the decoder that prevents the predictions and intermediate states to attend to positions equal to or greater than the current one, which means that the internal state won't change even if you recomputed them at every decoding step.

Update: of course it is technically possible to recompute past tokens attending future tokens, but then, what do you do with the future tokens after you re-compute the past ones? Do you recompute them? This is a totally different beast, which has been studied to some degree and is referred to as "iterative refinement". An example can be found in article "Deterministic Non-Autoregressive Neural Sequence Modeling by Iterative Refinement". AFAIK, this kind of approach has not studied in autoregressive models, only in non-autoregressive ones.