Inference Process in Autoregressive Transformer Architecture

Question

I'm abit confused about how the inference/prediction process works in transformer.

For simplicity suppose it is a transformer for translation. My understanding is that when training, the whole input and output sentences are put into model. This is possible because of the causal mask in the decoder that prevents the model from cheating and looking ahead.

Then, once the weights have been trained, the inference/prediction works by placing </s> or start sentence tag in the beginning with padding. The predicted word is then concatenated until </s> is the predicted word. My confusion arises from how the predicted word is acquired.

The causal mask ensures that the first predicted token (X_1 below) is only a function of the first token (i.e. is not affected by the padding we used in the other tokens. So our first predicted word/token should be taken from the first, and subsequently once we concatenated k words it should be taken from k+1 th output position. See the diagram below for clarity.

However, I've been using nlp.seas.harvard.edu/annotated-transformer/ as reference (and also checked another tensorflow tutorial), and they seem to take the predicted word/token as the last token (i.e. X_N).For example, under the inference section of the link above, we have:

prob = test_model.generator(out[:, -1])         
_, next_word = torch.max(prob, dim=1)         
next_word = next_word.data[0]         
ys = torch.cat([ys, torch.empty(1, 1).type_as(src.data).fill_(next_word)], dim=1)

Thus, my question is whether I'm misunderstanding the model of misunderstanding the code?

Answer 1

I think that your misunderstanding comes from this:

Then, once the weights have been trained, the inference/prediction works by placing or start sentence tag in the beginning with padding.

There is no padding at all. In the first iteration, the length of the input is only one token (i.e. for </s>). Once the prediction for the first token is computed, it is appended to the input, getting a tensor of length 2, which is then used as input, and so on. At each step, the length of the input increases by 1. Therefore, the prediction for the current step is always the last predicted token.