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!