[TTS] Add cosine distance option to TTS aligner

[rlangman]

What does this PR do ?

Add configuration to TTS aligner which uses scaled cosine distance instead of euclidean distance.

Collection: [TTS]

Changelog

• Add code for cosine distance based alignment

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

Additional Information

The current aligner used for FastPitch training has a few problem:

1. It aligns slowly for datasets with lots of speakers.
2. If binarization loss is enabled/warmup too soon then alignment learning stops before accurate alignments are learned.
3. When fine-tuning, the model does not learn alignment of target speaker well.

Using cosine distance instead of L2 distance can fix these problems. Cosine distance in general a lot more accurate and stable when it comes to similarity problems. It also avoids possible issues with the binarization loss because the scale of the cosine distance is fixed regardless of the scale or content of the alignment embeddings.

To test this I trained FastPitch on the VCTK dataset for 300k steps and provide the alignments and loss graphs below.

As seen, using cosine distance aligns faster, produces more accurate and sharper alignments, has a lot lower alignment loss (especially binarization loss), and as a result converges to a model with significantly better accuracy all around (mel, duration, pitch, energy). Though the gap between the two would likely shrink if we train it for several million steps.

Alignments

Example 1

L2

Cosine

Example 2

L2

Cosine

Alignment Losses

Reconstruction Losses

[github-actions]

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

[github-actions]

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

[XuesongYang]

LGTM. Added some neat-picks.

One discussion: I observed many @staticmethod member functions defined in class AlignmentEncoder. I wonder if nemo.collections.tts.parts.utils.helpers is a better place to hold them all?