Ignore whisper previous tokens during loss calculation

[DavraYoung]

Feature request

I am training whisper using previous prompts and according to whitepaper, its better to ignore the loss for previous tokens.

We only mask out the training loss over the previous context text, and train the model to predict all other tokens

Which means:
For: <|startofprev|>Here goes previous transcription tokens<|startoftranscript|><|transcribe|><|en|>hello world!<|endoftext|>
the loss must be masked for all tokens that are going before<|startoftranscript|>.

Simply putting -100 in labels is not correct since model actually attend to the decoder_input_ids.
Training without masking gives this:

1. Start training whisper with prompts (so called "previous tokens")
2. Observe whisper learning to predict prompts. (which its not supposed to do)
3. Observe loss never dropping below 1.0
4. Observe a lot of hallucinations

• transformers version: 4.34.1
• Platform: Windows-10-10.0.22621-SP0
• Python version: 3.9.13
• Huggingface_hub version: 0.16.4
• Safetensors version: 0.3.1
• Accelerate version: 0.21.0
• Accelerate config: not found
• PyTorch version (GPU?): 2.1.0.dev20230803+cu121 (True)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: no

cc: @sanchit-gandhi

Motivation

Finetuning whisper on audio with text context.

Your contribution

class ExtendedWhisperForConditionalGeneration(WhisperForConditionalGeneration):
    def __init__(self, config: WhisperConfig):
        super().__init__(config)
        self.model = WhisperModel(config)
        self.post_init()

    def forward(
        self,
        input_features: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.LongTensor] = None,
        decoder_input_ids: Optional[torch.LongTensor] = None,
        decoder_attention_mask: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        decoder_head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        decoder_inputs_embeds: Optional[Tuple[torch.FloatTensor]] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], Seq2SeqLMOutput]:
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        if labels is not None:
            if decoder_input_ids is None and decoder_inputs_embeds is None:
                decoder_input_ids = shift_tokens_right(
                    labels, self.config.pad_token_id, self.config.decoder_start_token_id
                )

        outputs = self.model(
            input_features,
            attention_mask=attention_mask,
            decoder_input_ids=decoder_input_ids,
            encoder_outputs=encoder_outputs,
            decoder_attention_mask=decoder_attention_mask,
            head_mask=head_mask,
            decoder_head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            decoder_inputs_embeds=decoder_inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        lm_logits = self.proj_out(outputs[0])

        loss = None
        if labels is not None:
            # Find the position of the <|startoftranscript|> token in labels
            transcribe_positions = (labels == 50258).nonzero(as_tuple=True)[1] - 1
            max_position = labels.shape[1]
            mask = (
                torch.arange(max_position)
                .expand(len(labels), max_position)
                .to(labels.device)
            )

            if (
                len(transcribe_positions) > 0
            ):  # Ensure there's at least one <|startoftranscript|> token
                mask = mask > transcribe_positions[:, None]

            # Modify the labels to be -100 (ignored in CrossEntropyLoss) for positions in the mask
            labels = torch.where(mask, labels, torch.tensor(-100).to(labels.device))

            loss_fct = CrossEntropyLoss()
            # Compute loss with modified labels
            loss = loss_fct(
                lm_logits.view(-1, self.config.vocab_size), labels.reshape(-1)
            )
        if not return_dict:
            output = (lm_logits,) + outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return Seq2SeqLMOutput(
            loss=loss,
            logits=lm_logits,
            past_key_values=outputs.past_key_values,
            decoder_hidden_states=outputs.decoder_hidden_states,
            decoder_attentions=outputs.decoder_attentions,
            cross_attentions=outputs.cross_attentions,
            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
            encoder_hidden_states=outputs.encoder_hidden_states,
            encoder_attentions=outputs.encoder_attentions,
        )

[DavraYoung]

cc: @sanchit-gandhi

[sanchit-gandhi]

I think the cleanest way of doing this is actually in the data collator, as per: https://github.com/huggingface/distil-whisper/blob/e35cbff09e04f6e710c72f33e0441d6bebc47ec1/training/run_distillation.py#L400

[DavraYoung]

@sanchit-gandhi I think data collator may not be suitable here, since we are dealing with loss masking.

As I understood to achieve this in data collator I will need to replace decoder_input_ids with -100 values, which means that previous prompt tokens will simply be removed from training process, which makes them useless.

In other words, If we simply replace initial prompt tokens in decoder_input_ids with -100 values, model will not be able to access the information in those tokens.

So, the decoder_input_ids should contain previous_prompt tokens, but loss on such tokens should be ignored at loss calculation step, but due to hf transformers whisper implementation, there is no way to mask the loss at the loss calculation step on certain token positions.

[DavraYoung]

oh, wait, sorry, I got it. You mean to supply decoder_input_ids and labels together? The decoder_inputs_ids will be shifted right labels and labels will have -100 on previous prompt positions

[sanchit-gandhi]

What we do is:

1. Shift the decoder input ids to the right to get the labels (targets)
2. Shift the input ids attention mask to the right to get the labels mask
3. Mask out where we have padded tokens in the labels (set to -100 to ignore in the loss)
4. Mask out the prompt in the labels (set to -100 to ignore the loss)

=> this means that the prompt is included in the decoder input ids, so the model can condition on the prompt, but is masked from the loss, so the model is not trained to predict the prompt. This makes sense intuitively: we don't want to train the model on the prompt, only the target transcription. Indeed, if you refer to page 3 of the Whisper paper, you'll see this is one-to-one the same as how they do it:

[DavraYoung]

Thanks for explanation.
Btw, what about task tokens? Like <|translate|>, <|transcribe|> etc.
Should they be ignored during loss calculation? It seems counterintuitive that the model predicts them.
For language and silence detection loss should definately be computed, but for task tokens its not clear.

During training those tokens impact loss calculation, since the model predicts them, which makes loss charts be less informative

For example:
<|startoftranscript|><|translate|> token1 token2 <|endoftext|>
during training step will give loss for each of the token being predicted:
For <|startoftranscript|> the loss will be calculated based on the probability of <|translate|> token, and if I have mostly trained on transcribe task, the loss will be slightly higher then it actually is, because the task token will be predicted incorrectly.

[sanchit-gandhi]

They should be counted, e.g. for the language token it corresponds to the model's prediction for the source audio language. If we mask these tokens, the model will not be trained to differentiate between translate/transcribe or tell what language the audio is in. At inference time, we 'force' the task tokens (e.g. task="transcribe") which sets the task token to our desired value -> this then controls what kind of generation we perform, so the model is forced to perform the task we specify.

[DavraYoung]

If the model's processor is anyway forcing these tokens(specifically transcribe and translate), and the model is never choosing by itself, then maybe it we can ignore calculating the loss on model prediction of such tokens during training?

By not including these tokens in loss calculation, the model can potentially focus more on learning the core aspects of the tasks (like transcribing content accurately) instead of predicting task tokens that are already predetermined.

For example:

In my current scenario the whisper model chooses the next step given the list of possible steps and the audio.
For example:
Actor says: "Remove the 231 todo from the list".

previous tokens:

1. Add todo item
2. Remove todo item
3. List todo items

learnable_tokens:

<|startoftranscript|><|predict_next_action|>
2. 231
<|endoftext|>

The previous tokens tokens are masked and the model still pays attention to them,
which, I think, means: masking task tokens will also not affect the model accuracy of performing a certain task

[DavraYoung]

for any one who faces same problem, adjust your data collator to supply both labels and decoder_input_ids.
Example: https://github.com/huggingface/distil-whisper/blob/e35cbff09e04f6e710c72f33e0441d6bebc47ec1/training/run_distillation.py#L400

[sanchit-gandhi]

We need to train on task/language since the model needs to learn how these tokens affect the output:

• task="translate" means speech translation, task="transcribe" means speech transcription. If we mask the task token, then the model has no way of differing between these tasks during training time, and so will perform them randomly at inference time
• Same for language (need to train on the language token so the model learns to differentiate between languages)

So what we do is train the model on these tokens during training time (model learns that task="translate" means speech translation), and then force them at inference time to control the behaviour that we want.

However, if you are only doing one task (resp. one language), you can mask the task (resp. language) token accordingly. This is what's done for English-only training in the original model for example: the tokens are removed entirely.