ProtT5 model training loss function

[Alex2975]

Dear Authors,

Thank you very much for the great work. I got a question and would appreciate your insights.

For the ProtT5 training, since it will predict the full sequence, not just the masked tokens. What is the loss function for the ProtT5 training, is it Torch crossentropyloss with reduction=SUM, or it is Torch crossentropyloss with reduction=MEAN?

[mheinzinger]

I would always recommend to use mean as you want the loss to be independent of the number of tokens you have in your batch

[Alex2975]

Thank you for the insights, @mheinzinger . If I want to compute the perplexity of a protein sequence from the ProtT5 model, how do I do it? Since it is using MLM objective, I think the following will not work. Could you please share some insights?

inputs = tokenizer(sequence, return_tensors="pt")["input_ids"]
labels = tokenizer(sequence, return_tensors="pt")["input_ids"]

Get outputs

outputs = model(**inputs, labels=labels)
loss = outputs.loss

Compute perplexity

perplexity = torch.exp(loss).item()
print(f"Perplexity: {perplexity}")

[mheinzinger]

Indeed, the above won't work. Especially, as perplexity is ill-defined for models trained via MLM.
I guess the best you can get is the pseudo perplexity where one masks one token at a time, reconstructs it, compute loss against groundtruth, and repeat for the full sequence before averaging & taking exponent.
We have some implementation for step-wise masking of ProtT5 here, maybe this helps.

[Alex2975]

Thank you so much for the tips, @mheinzinger . I found how the "masks one token at a time" from the link you shared. Will try that.

[mheinzinger]

Just in case; we now have an example for continuing pre-training.