1
$\begingroup$

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.

My understanding of Training Process

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.

My understanding of Inference Process

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?

$\endgroup$

1 Answer 1

1
$\begingroup$

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.

$\endgroup$
2
  • $\begingroup$ This account is still new so my upvote doesn't show, but thanks! For some reason I had it in my head that decoder input size has to be fixed. $\endgroup$ Commented Sep 19, 2022 at 12:23
  • $\begingroup$ I'm glad my answer helped. Please consider marking the answer as correct with the tick on the left. $\endgroup$
    – noe
    Commented Sep 19, 2022 at 13:04

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.