I'm trying to do SMILES chemical representation prediction from a large dataset (Around 5M Samples) to teach it do predict another downstream task. The model's part responsible for generating the data is a decoder embedding layer that roughly looks like this:
self.decoder_embedding = nn.Embedding(len(tokenizer), hidden_size)
decoder_layer = nn.TransformerDecoderLayer(
d_model=hidden_size,
nhead=heads,
dim_feedforward=hidden_size,
dropout=dropout,
batch_first=True,
norm_first=True
)
self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=2)
self.smiles_generation_head = nn.Linear(hidden_size, 125)
The inputs to the model includes smiles_tokens which have been randomly masked of shape [batch_size, 125] and a padding attention mask. The output is simply pretext_predictions of shape [batch_size, 125].
def mask_tokens(inputs, tokenizer, mask_prob=0.15):
masked_inputs = inputs.clone()
mask_token_id = tokenizer.mask_token_id
# Mask a random selection of tokens
random_mask = torch.rand(masked_inputs.shape) < mask_prob
masked_inputs[random_mask] = mask_token_id
return masked_inputs
masked_smiles_tokens = mask_tokens(non_share_smiles_tokens, tokenizer)
smiles_tokens = model(masked_smiles_tokens, *features)
When this is passed into the loss function, it starts at an incredibly high loss value (roughly 10K and above). When the batch size is increased, the loss value also doubles and sometimes triples to around 30K or above.
smiles_tokens = nn.ConstantPad1d(
(0, 125 - smiles_tokens.shape[1]),
0
)(smiles_tokens).float()
padding_mask = smiles_tokens != 0
loss = loss_fn(
pretext_predictions[padding_mask],
smiles_tokens[padding_mask],
)
This is the loss function definition:
def loss_fn(inputs, targets):
ce_criterion = nn.CrossEntropyLoss(reduction='mean')
ce_loss = ce_criterion(inputs, targets)
return ce_loss
What is causing this high loss value? I tired normalizing my input features apart from the SMILES token indices and that doesn't seem to solve the issue. I noticed that when I create random tensors like this:
import torch
import torch.nn as nn
import random
# Example tensors (replace with your actual data)
predictions = torch.randn(16, 50).softmax(dim=1)
targets = torch.randint(0, 100, (16, 50)).float()
mask = torch.randint(0, 2, (16, 50))
# Create the loss function
loss_fn = nn.CrossEntropyLoss()
# Calculate the loss
loss = loss_fn(predictions[mask], targets[mask])
print(loss)
It also result in huge loss values.
tensor(9981.9561)
My ultimate aim is a BCE task and the loss from BCE is very small (0 to 1) compared to the CE loss. This makes it difficult to asses whether the model is making any progress. What should I do? Do I simply just find another loss function? NLLLoss also doesn't seem to being doing well but CE Loss uses that under the hood so I'm guessing it stems from that?
Note: pretext_predictions are logits, and smiles_tokens are indices for the tokenizer vocabulary of size 37. The learning rate is 1e-3 using Adam and model size is only 5M parameters.
