I'm developing an encoder-decoder based transformer model and I would like to ask if there are ways to incentivize or penalize certain tokens during training.

I'm working on a translation task where the encoder input must be decoded into its proper product name. I have labels such as brand, name, and unit of measure, etc which are available during training but not on inference.

Currently when predicting the brand portion (which usually appears early in the sequence) of the output, the heatmap shows that it does not give focus to the latter part of the encoder which produce an output that the brand and product name, and unit of measure does not belong to each other.

I was thinking if there's a way to force the transformer during training to give more weight to different token types other that its own.

For example:

  1. Brand tokens (decoder) should give more weight to name tokens (encoder) than other brand tokens (encoder)
  2. Name tokens (decoder) should give more to brand token (encoder) and unit of measure token (encoder)

1 Answer 1


Yes, that is completely possible. Instead of using a mean of loss weights, you can use higher weights for tokens of choice and calculate weighted mean.

Pytorch code for the same Code Source

log_prob = torch.tensor([[-0.0141, -4.2669],
                     [-0.0141, -4.2669]])
target = torch.tensor([0, 1])
weight = torch.tensor([2.0, 3.0])
criterion = nn.NLLLoss()
criterion_weighted = nn.NLLLoss(weight=weight)

print(criterion(log_prob, target))
> tensor(2.1405)
print(criterion_weighted(log_prob, target))
> tensor(2.5658)

The weights will enforce/bias your transformer's attention mechanism to focus on specific relationships in the data. Note the weights will be a hyper parameter you need to fine tune.


You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .