(air whooshing)
(logo popping)
(metal sliding)
- In this video, we take a look at setting up
a custom loss function for training.
In the default loss function, all samples,
such as these code snippets, are treated the same
irrespective of their content but there are scenarios
where it could make sense to weight the samples differently.
If, for example, one sample contains a lot of tokens
that are of interest to us
or if a sample has a favorable diversity of tokens.
We can also implement other heuristics
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 create a weighted sum
or average over all samples
to get the final loss for the batch.
Let's have a look at a specific example.
We want to set up a language model
that helps us autocomplete 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 model's input and predicted logits,
as well as the key tokens, as input.
First, the inputs and logits are aligned.
Then the loss per sample is calculated,
followed by the weights.
Finally, the loss and the weights are combined and returned.
This is a pretty big function, so let's take a closer look
at the loss and the 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 the matrix
with a row for each sample in the batch and a column
for each position in the sequence of the sample.
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.
We only want to know
how many times keywords occurred per sample,
so we can sum overall keywords and all positions per sample.
Now we're almost there.
We only need to combine the loss with the weight per sample.
We do this with element wise multiplication
and then average overall 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 model to get the logits
and then we can 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 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 Accelerate.
(air whooshing)