[Bug] BCELoss should not be masked

[iamanigeeit]

I have trained Tacotron2 but during eval / inference, it often doesn't know when to stop decoding. This is a known issue in seq2seq models and i was trying to solve it in TensorFlowTTS when i gave up due to Tensorflow problems.

Training with enable_bos_eos=True helps a bit but the output is still 3x the ground truth mel length for shorter audio: see length_data_eos.csv vs length_data_no_eos.csv

One reason is the BCELossMasked criterion -- in its current form, it encourages the model never to stop decoding once it has passed mel_length. Some of the loss results don't quite make sense, as seen below:

import torch
def sequence_mask(sequence_length, max_len=None):
    if max_len is None:
        max_len = sequence_length.data.max()
    seq_range = torch.arange(max_len, dtype=sequence_length.dtype, device=sequence_length.device)
    # B x T_max
    mask = seq_range.unsqueeze(0) < sequence_length.unsqueeze(1)
    return mask

from torch.nn import functional
length = torch.tensor([95])
mask = sequence_mask(length, 100)
pos_weight = torch.tensor([5.0])
target = 1. - sequence_mask(length - 1, 100).float()  # [0, 0, .... 1, 1] where the first 1 is the last mel frame
true_x = target * 200 - 100  # creates logits of [-100, -100, ... 100, 100] corresponding to target
zero_x = torch.zeros(target.shape) - 100.  # simulate logits if it never stops decoding
early_x = -200. * sequence_mask(length - 3, 100).float() + 100.  # simulate logits on early stopping
late_x = -200. * sequence_mask(length + 1, 100).float() + 100.  # simulate logits on late stopping

# if we mask
>>> functional.binary_cross_entropy_with_logits(mask * true_x, mask * target, pos_weight=pos_weight, reduction='sum')
tensor(3.4657)  # Should be zero! It's not zero because of trailing zeros in the mask
>>> functional.binary_cross_entropy_with_logits(mask * zero_x, mask * target, pos_weight=pos_weight, reduction='sum')
tensor(503.4657)
>>> functional.binary_cross_entropy_with_logits(mask * late_x, mask * target, pos_weight=pos_weight, reduction='sum')
tensor(503.4657)  # Stopping late should be better than not stopping at all. Again due to trailing zeros in the mask
>>> functional.binary_cross_entropy_with_logits(mask * early_x, mask * target, pos_weight=pos_weight, reduction='sum')
tensor(203.4657)  # Early stopping should be worse than late stopping because the audio will be cut

# if we don't mask
>>> functional.binary_cross_entropy_with_logits(true_x, target, pos_weight=pos_weight, reduction='sum')
tensor(0.)  # correct
>>> functional.binary_cross_entropy_with_logits(zero_x, target, pos_weight=pos_weight, reduction='sum')
tensor(3000.)  # correct
>>> functional.binary_cross_entropy_with_logits(late_x, target, pos_weight=pos_weight, reduction='sum')
tensor(1000.)
>>> functional.binary_cross_entropy_with_logits(early_x, target, pos_weight=pos_weight, reduction='sum')
tensor(200.)  # still wrong

# pos_weight should be < 1 to penalize early stopping
>>> functional.binary_cross_entropy_with_logits(zero_x, target, pos_weight=torch.tensor([0.2]), reduction='sum')
tensor(120.0000)
>>> functional.binary_cross_entropy_with_logits(late_x, target, pos_weight=torch.tensor([0.2]), reduction='sum')
tensor(40.0000)
>>> functional.binary_cross_entropy_with_logits(early_x, target, pos_weight=torch.tensor([0.2]), reduction='sum')
tensor(200.)  # correct

For now i am passing length=None to avoid the mask and setting pos_weight=0.2 to experiment. Will update the training results.

Additional context

I would also propose renaming stop_tokens to either stop_probs or stop_logits depending on context. Currently, inference() produces stop_tokens that represent stop probabilities, while forward() produces the logits before sigmoid. Confusingly, both are called stop_tokens.

[erogol]

Nice catch also thanks for sharing the code to reproduce. It made really easy to test out your arguments.

Masking really seems like breaking the loss computation.

If you want you can send a ✨PR✨

We could also add the cases above in the tests.

Changing stop_tokens to something else is not something I'd like to even though I agree with you since it is referred in the docs and discussions.

[iamanigeeit]

Nice catch also thanks for sharing the code to reproduce. It made really easy to test out your arguments.

Masking really seems like breaking the loss computation.

If you want you can send a sparklesPRsparkles

We could also add the cases above in the tests.

Changing stop_tokens to something else is not something I'd like to even though I agree with you since it is referred in the docs and discussions.

I have rerun training on LJSpeech without masking, with pos_weight=0.2 and with an added stopnet_alpha=100.0 (since reducing pos_weight reduces the stopnet loss value). It doesn't reduce the output_lengths and in fact makes it even longer: length_data_maskedloss.csv vs length_data_nomaskedloss.csv

BUT importantly, it solved one major problem. Every sample output i have checked using masked loss suffers from repetition (a well-known problem in NLP):

LJ025-0110

Using loss without masking:

The resulting silence can later be trimmed out, so it is ok, it just makes decoding slower.

[erogol]

And for how many steps you've trained the model and with what config?

Maybe longer training would make it better.

[iamanigeeit]

I used the same config as recipes/ljspeech/tacotron2-DCA/train_tacotron_dca.py except batch_size=32 (due to GPU memory limit) and r=1 (i think r=1 is the correct one for Tacotron2). Training was for 100k steps each.

[lexkoro]

r=1 is rly hard to train using tacotron2, especially if the dataset has samples with long silences.

Also you might want to try activating dropout during inference, I remember it improved my inference by a lot when using dca.

[erogol]

@iamanigeeit thanks for sharing your updates.

Do you intend to send a ✨PR✨ that fixes the BCELoss?

Or we can also handle it but I guess you deserve the attribution for this. That is why it'd be nice if you could send the PR.

[erogol]

r=1 is rly hard to train using tacotron2, especially if the dataset has samples with long silences.

Also you might want to try activating dropout during inference, I remember it improved my inference by a lot when using dca.

Finding the right DB value for trimming the silence would help. I agree that training with r=1 is a hard task when the dataset is not processed well for TTS.

[iamanigeeit]

@iamanigeeit thanks for sharing your updates.

Do you intend to send a sparklesPRsparkles that fixes the BCELoss?

Or we can also handle it but I guess you deserve the attribution for this. That is why it'd be nice if you could send the PR.

Haha, i don't need the attribution... but i need help using Git as i haven't used it on multi-person projects before. I now have a version cloned from the 33aa27 commit with all my commits pushed to my own repo, including commits on other files. How do i fix this?

Here are the files to update: updates.zip

[Edresson]

Follow some steps:

1st Fork the 🐸 TTS repository (use the button "fork" at the top of the page)

2st Clone from your Fork (dev branch). The command will be some like:
git clone https://github.com/iamanigeeit/TTS.git -b dev

3st Change the files that you need.

4st Commit the changes with the commands (obs: Change the commit message :)):

git add .
git commit -m "Commit message"

5st Push the commits to your fork with the command: git push

6st Go to your fork (https://github.com/iamanigeeit/TTS). Github will identify that you've made changes and suggest the pull request and it will show a pull request button below "Go to file", "Add file" and "code" buttons. Now you can click on the pull request button and send a pull request from your dev branch to Coqui's dev branch :).

[iamanigeeit]

@erogol I think losses should not be masked at all.

From the original Tacotron paper:

It’s a common practice to train sequence models with a loss mask, which masks loss on zero-padded frames. However, we found that models trained this way don’t know when to stop emitting outputs, causing repeated sounds towards the end. One simple trick to get around this problem is to also reconstruct the zero-padded frames.

This is what i've found as well.

[iamanigeeit]

@Edresson thanks for the suggestion -- i've made a lot of other changes in prior commits to my own repo, so probably i will create a new fork and update from there...

[erogol]

I've tried no masking but masking worked better in all of my experiments although it was buggy.

[iamanigeeit]

How did you manage to solve the non-terminating outputs problem?

[erogol]

I've not observed this problem with my models after training long enough.

[iamanigeeit]

Thanks, i was thinking 100k steps on batch size 32 would be enough, but it seems not... you haven't encountered inference frequently hitting max_decoder_steps?

[erogol]

you can change max_decoder_steps in the model config.

[iamanigeeit]

I set max_decoder_steps to a big number to test whether the model stops. I think it is a problem when the model doesn't stop generating, although limiting the decoder steps is a workaround.

[erogol]

Sorry I misunderstood your question.

Thanks, i was thinking 100k steps on batch size 32 would be enough, but it seems not... you haven't encountered inference frequently hitting max_decoder_steps?

Yes, I've not observed the model hitting the max_decoder_steps after enough training.

[erogol]

@Edresson can you look at this when you have time?