Training a Speaker Embedding from Scratch with Triplet Learning

For the past few months, I have been researching on building an end to end speaker identification system with deep learning. The area of research I focused on is metric learning. Knowing almost nothing going in, I have significantly underestimated the effort required (surprise!). While the effort is ongoing, I have learnt much along and have some preliminary results I am excited to share.

In this blog post, I will explain what metric learning is, what is speaker identification, how metric learning applies in this context, and share the lessons I learnt from applying those knowledge to build a speaker embedding function.

tl;dr I share my learnings on how I built a speaker embedding function with deep learning in this blog post

A bit of background on metric Learning

Metric learning is about learning a function that translates an input such as an image to another hyperspace, commonly referred to as latent space, where one can compare different instances on some metrics, usually semantic similarity. That is, metric learning is not about learning to classify an image to a predetermined class; rather, metric learning is about learning a similarity function that we can use to compare how similar two inputs are. Why would we want to such a function? This is a useful technique when the input space is large and it’s not feasible to have a sample from every class to train a classifier.

Perhaps the most well known application of metric learning are word vectors. A word vector function takes words as inputs and output embeddings in a vector space where words are clustered around semantic similarity. For example, in a well trained word embedding, “queen” and “woman” will be closer together than “king” and “woman” as the training implicitly learns the concept of gender. Another interesting application of metric learning is finding similar or duplicate images. FaceNet [1], which does person identification using face images, is an example of such application.

Speaker identification problem in detail

Now that we covered the basics of metrics learning with examples such as word vectors and faceNet, let’s look at the speaker recognition problem in more detail.

In the research literature, speaker identification is divided into 2 categories. If we are identifying a speaker using a fixed phrase such as “Ok Google” or “Alexa”, this is referred to as text-dependent speaker identification. If we are identifying a speaker using only the voice characteristics without relying on specific words, this is known as text independent speaker identification. For instance, a well learnt text independent embedding should be able to identify my voice whether I say the words “Hello World” or “Goodbye World”. I decided to focus on text independent identification from the beginning since it solves the speaker identification problem in the general settings. See Google’s paper on building “OK Google” text dependent speaker identification at [2].

In summary, speaker identification in this post means recognizing the speaker identity based on a sample of a speaker’s voice, independent of the spoken words. To be able to accomplish this, we need to train a function that transform a speaker voice to an embedding space where similar voices are clustered near each other and different voices are far away from each other. With this function, we will then be able to identity speakers based on its neighbors in the embedding hyperspace.

Quick thoughts on data processing and dev environments

With the problem defined, the next step to tackle is data and development environment. You probably heard before that machine learning problem is 90% data problem and 10% machine learning. You may also have thought this was an exaggeration. I did and I was wrong. In the context of engineering effort for this experiment, I spent a significant amount of time setting up scalable data processing pipelines and repeatable machine learning experiments. This effort itself is worth a separate blog post so I won’t belabor the effort further. I ended up going with pytorch as the deep learning framework, chose Mel-frequency cepstral coefficients (MFCCS) as the data representation, used HDF5 with data versioning, recorded my experiment parameters in a yaml config, and used tensorboard as the visualization tool.

Along the same vein, finding quality datasets for this experiment took nontrivial efforts as well. According to Deep Speaker [3] from Baidu, the models they trained relied on a dataset with one million plus speakers. As far as I know, there is no public dataset with nearly as many speakers. With that in mind, I started gathering my own dataset and currently have about 10,000 speakers across multiple languages; however, for the experiments and visualizations I share in this blog post, I used publicly available dataset comprising of TIMIT, VCTK, and dev and test LibriSpeech dataset. These dataset include approximately 880 English speakers, which I split into 680 training speakers and 200 test speakers. Note that there are other public speaker dataset such as Mozilla’s common voice and VoxCeleb. I have not had a chance a to explore them in detail yet but I plan to do so later.

From Contrastive Loss to Triplet Loss

With the problem and dataset defined, let’s look in more detail how we can utilize metric learning to tackle the speaker recognition problem. Typically, we will utilize a siamese network architecture. This means we will create a deep neural network with a fully connected layer as the output layer and train the network to minimize the difference between voice embeddings from the same speaker and maximize the difference between voice embeddings from different speakers. For a good introduction to siamese architecture, see [4].

As shown in fig 0 above, siamese network is an architecture and does not prescribe what each layer of the neural network encompasses. What is important is that two paths of the network share the same weights to make output embedding directly comparable. With a neural network architecture, the next step is to train the network with a suitable loss function. In the metric learning space, there are 2 common loss functions: contrastive loss and triplet loss.

Training with contrastive loss involves taking a batch of sample pairs as inputs to train the network to separate different class embeddings with a minimum distance called margin while trying to minimize the distance between embeddings of the same class. The distance function is not prescribed but is usually euclidean or cosine. During training, one trains over a generated list of pairs such as (x, y) where x might be from the same class as y half the time. Contrastive loss is defined in fig 1 below. For an example of contrastive loss in the literature, see [5].

On the other hand, triplet loss involves 3 inputs, commonly called anchor, positive and negative sample. Anchor is an arbitrary sample and positive is another sample from the same class as anchor while negative is a sample which has a different class than either anchor or positive.

Instead of focusing on distance between positive pairs or negative pairs, triplet loss encourages the network to maximize the margin of distance between the distance of anchor and positive versus the distance between anchor and negative. That is, it encourages the anchor and positive to be as close as possible while pushing anchor and negative to be as far as possible in the embedding hyperspace. Triplet loss pseudo code is presented in fig 3 below. See [6] for an introduction on triplet learning.

Finally, as an aside, it goes without saying that are myriads of other techniques in this research space. I decided to pursue contrastive loss and triplet loss due to their conceptual simplicity and existing literatures’ application on similar problems.

From theory to practice and back again

While I was able to successfully train an embedding function with contrastive loss, the results were ok at best so I decided to focus on triplet loss as the primary technique. On an intuitive level, triplet loss optimizes for a margin on inter distances between speakers so it should produce more discriminative embeddings in theory.

This is where I experienced the greatest difficulty in this project. While many of the papers I read present great success with the technique, implementing triplet training efficiently and correctly proved to be frustrating and error-prone.

First, generating triplets is an O(n³) operation — in other words, it is infeasible to generate all possible triplets for training. If we cannot generate all triplets for training, maybe we can just randomly generate a subset of triplets instead?

Unfortunately, triplet loss is extremely sensitive to the training samples and uniform sampling usually leads to failed training. This is the second peril when it comes to triplet loss — training often results in plateaued training loss or local minima that the network fails to escape from. One of the most common failure scenario I have repeatedly observed is that the network converges to a constant output (ie it outputs the same embedding for any input) and fails to escape the local minima of loss equalling margin. For an excellent overview on triplet learning and some common issues, see [7].

There are also many other papers that focus on good sampling strategy for triplets to ensure training succeeds. Most of them center around the idea of hard negative mining — that is, we need to generate triplets that violate the margin requirements to provide meaningful signals to help the network learn. If most of the triplets during training produce zero loss, the network will not be able to meaningfully decrease its loss. See [8] for an example on sampling strategy to overcome the difficulty of triplet training regime.

While the research suggests hard negative mining and its variants work, they are computationally expensive as we have to generate new triplets, forward them through the network to get their embeddings, and pick out triplets that violates the margin requirements. These steps must be performed after each training update, resulting in extremely slow training times.

All of these problems led me back to square one where I scoured the literature again to try to identify a new direction. This is when I came across [9]. This was the light-bulb moment for me in this project. These papers suggested instead of generating triplets, we learn an embedding for each class (speaker in my case) and use the learnt embedding as a proxy for triplets as part of the training. In other words, we can train end to end without the computationally expensive step of resampling triplets after each network update.

Proxy based triplet learning is both easier to implement and quicker to train. To verify that my implementation works, I applied the method to learn an embedding on everyone’s favorite dataset — MNIST. You can see the results in the following 2 figures.

As show in the 2D and 3D T-SNE in the figures above, the method produces well separated clusters for each digits.

Excited by this validation, I applied the triplet proxy learning to the 680 speakers dataset and visualized the result on 50 test speakers via T-SNE. Despite the relatively small training dataset, you can see clusters forming around speaker identities as demonstrated in the animation below.

While the embedding is far from perfect, it validates the basic strategy of identifying speakers via its neighbors is viable. Here is the same embedding dataset viewed in 3D.

Triplet of Lessons Learnt

In conclusion, I like to highlight three major lessons I learnt. First, I cannot stress how important it is to get the tooling right from the very beginning. Without a fast iterative environment, everything down the road becomes exponentially more difficult. Invest in making your development environment fast and repeatable.

Second, it is important to review the literature thoroughly. As I shared in this post, there are many different ways of implementing triplet training regime and it is difficult to know whether a particular technique will work for your problem. So try as many of them as possible and go back to the literature if needed.

Finally, start extremely small and verify incrementally. One of my big mistake early on in this project is not sanity check with MNIST, whether it was to check tooling correctness or verify a training technique works. By the time I realized my mistake, I had already wasted days of GPU hours and human time debugging issues that would have been caught by starting with MNIST.

Thanks for making this far! If I did something right, maybe you learnt something worthwhile from my journey on exploring and taking deep learning from theory to practice.

Next Step

This project is far from complete! There are many exciting extensions such as training on multi language speaker dataset, identifying accents, making embedding function noise tolerant, etc. If anything, this is just the beginning of the journey.

References

[1] https://arxiv.org/abs/1503.03832
[2] https://research.google.com/pubs/pub44681.html
[3] https://arxiv.org/abs/1705.02304
[4] https://cs.nyu.edu/~sumit/research/assets/cvpr05.pdf
[5] http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf
[6] https://arxiv.org/abs/1412.6622
[7] https://arxiv.org/abs/1703.07737
[8] https://arxiv.org/abs/1704.01285
[9] https://arxiv.org/abs/1703.07464