In this video we take a look at setting up a custom loss function for training. In the default loss functions all samples such as these code snippets are treated the same irrespective of their content, but there are scenarios where you it could make sense to weight the samples differently. If for example one sample contains a lot of tokens that or of interest to us or it has a favourable diversity of tokens. We can also think of other heuristics we can implement with pattern matching or other rules. For each sample we get a loss value during training and we can combine that loss with a weight. Then we can for example create a weighted sum to get the final loss for a batch. Let’s have a look at a specific example: we want to setup a language model that helps us autocomplete complete common data science code. For that task we would like to weight samples stronger where tokens related to the data science stack, such as pd or np, occur more frequently. Here you see a loss function that does exactly that for causal language modeling. It takes the models it takes the model’s inputs and predicted logits as well as the key tokens as input. First the inputs and logits are aligned, then the loss per sample is calculate followed by the weights. Finally the loss and weights are combined and returned. This is a pretty big function so let’s take a closer look at the loss and weight blocks. During the calculation of the standard loss the logits and labels are flattened over the batch. With the view we unflatten the tensor to get a matrix with a row for each sample in the batch and a column for each position in the sequence of the samples. We don’t need the loss per position so we average the loss over all positions for each sample. For the weights we use boolean logic to get a tensor with 1s where a keyword occurred and 0s where not. This tensor has an additional dimension as the loss tensor we just saw because we get the information for each keyword in a separate matrix. Only want to know how many times keywords occurred per sample so we can sum over all keywords and all positions per sample. Now we are almost there, we only need to combine the loss with the weight per sample. We do this with element wise multiplication and then average over all samples in the batch. In the end we have exactly one loss value for the whole batch. And this is the whole necessary logic to create a custom weighted loss. Let’s see how we can make use of that custom loss with Accelerate and the Trainer In Accelerate we just pass the input_ids to the models to get the logits and can then call the custom loss function. After that we continue with the normal training loop by for example calling backward. For the Trainer we can overwrite the compute loss function of the standard trainer. We just need to make sure that that we return the loss and the model outputs in the same format. With that you can integrate your own awesome loss function with both the trainer and accelerates.