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.

[rexxxx1234]

Hi @mheinzinger,

Thank you very much for the great work. I checked out the continuing pre-training script, and saw that the loss is only calculated on the masked token for training.

Just wondering for the original protT5 pretraining, is the loss calculated on the full sequence or just the masked token? Many thanks!

[mheinzinger]

Just wondering for the original protT5 pretraining, is the loss calculated on the full sequence or just the masked token? Many thanks!

For original ProtT5 pre-training we computed loss over all tokens (so also non-masked ones) but Ahmed switched this in Ankh pre-training to computing loss only for masked positions which worked equally well. For compute efficiency reasons (and also to avoid biasing loss towards copy&paste of non-masked input tokens), we decided to focus the loss only on the masked items.