How do Transformer decoders handle arbitrary length input?

Question

I am working through a Tensorflow Neural Machine Translation tutorial (https://www.tensorflow.org/text/tutorials/transformer) and am confused about how the decoder handles inputs when making inferences after it has been trained.

In the section where we create a class for translating a sentence (https://www.tensorflow.org/text/tutorials/transformer#run_inference) it appears that we feed the decoder an array populated with the START token only, then append the last element from the predicted sequence it makes to the growing translated sentence. This does not make sense to me as we trained the transformer on a fixed length sequence that was padded to length MAX_TOKENS=128, so I can't figure out why it would be able to accept input of an arbitrary length tensor.

Here is the code for inference in question that confuses me:

    # As the output language is English, initialize the output with the
    # English `[START]` token.
    start_end = self.tokenizers.en.tokenize([''])[0]
    start = start_end[0][tf.newaxis]
    end = start_end[1][tf.newaxis]

    # `tf.TensorArray` is required here (instead of a Python list), so that the
    # dynamic-loop can be traced by `tf.function`.
    output_array = tf.TensorArray(dtype=tf.int64, size=0, dynamic_size=True)
    output_array = output_array.write(0, start)

    for i in tf.range(max_length):
      output = tf.transpose(output_array.stack())
      predictions, _ = self.transformer([encoder_input, output], training=False)

      # Select the last token from the `seq_len` dimension.
      predictions = predictions[:, -1:, :]  # Shape `(batch_size, 1, vocab_size)`.

      predicted_id = tf.argmax(predictions, axis=-1)

      # Concatenate the `predicted_id` to the output which is given to the
      # decoder as its input.
      output_array = output_array.write(i+1, predicted_id[0])

      if predicted_id == end:
        break

self.transformer takes in the "encoder_input" which is fed to the encoder and "output" which is fed into the decoder, where we then take the last prediction in the sequence returned from the decoder, append it to "output", then repeat the process. But I don't understand how the decoder can process a tensor of length 1, then 2, then 3, and so on if it is always expecting a tensor of length 128.

It seems like the input to the decoder should be padded to 128 tokens and then the prediction at the i-th index should be appended to the output (rather than always the last element in the predicted sequence) before repeating.

Is the tutorial mistaken in the implementation on of how this transformer preforms inference? Or am I missing something that Tensorflow does behind the scenes that allow this to work?

Note, I realize the title to this question is similar to How can Transformers handle arbitrary length input? but the answer given appears to verify that the input to the decoder should be fixed length and padded, which contradicts how inference is done in the tutorial I am referencing. So I hope this question isn't considered a duplicate but I am asking a similar question in the context of the mentioned tutorial.

Thank you for any insight someone might be able to provide!

Answer 1

At training time, the decoder uses a so-called triangular mask that prevents the self-attention to "look" at positions that are yet to be decoded. Because of that, nothing prevents the decoder to use shorter sequences than those it was trained with. The padding would be masked out anyway.

A more interesting question is whether the decoder can work with longer sequences than it was trained for. The answer is: it depends. In the original Attention is all you need paper, analytically defined position encoding was used that should in theory generalize beyond the lengths it was trained for. This is also used in the tutorial the question refers to. However, in practice, learned position embeddings are used (e.g., BERT uses a fixed length of 512 tokens) that seem to deliver better results in practice. In such a case, the decoder simply cannot generalize beyond the maximum length, simply because it does not have embeddings for those positions.