Powerset multi-class cross entropy loss for neural speaker diarization

Introduction

Speaker diarization is the task of partitioning an audio stream into homogeneous temporal segments according to the identity of the speaker. Most dependable diarization approaches consist of a cascade of several steps : voice activity detection to discard non-speech regions, speaker embedding to obtain discriminative speaker representations, and clustering to group speech segments by speaker identity. This family of approaches has two main drawbacks:

errors made by each step are propagated to the subsequent steps, possibly escalating into even larger errors;

they are not designed to detect (let alone assign to the right speakers) overlapping speech regions – an additional post-processing step is needed to handle them .

To circumvent these limitations, a new family of approaches have recently emerged, rethinking speaker diarization completely . Dubbed end-to-end diarization (EEND), the main idea is to train a single neural network that ingests the audio recording and directly outputs the diarization, hence addressing the error propagation issue. They formulate the problem as a multi-label classification task, allowing multiple overlapping speakers to be active simultaneously. Furthermore, permutation-invariant training is the critical ingredient that turns clustering – an intrinsically unsupervised task – into a supervised classification task. EEND approaches do have a few significant limitations:

they struggle to predict the correct number of speakers, especially when it is larger in test data than during training – additional mechanisms such as encoder-decoder attractors were proposed to partially solve this issue ;

training such approaches is extremely data hungry as each conversation in the training set constitutes a single training sample – one has to resort to synthetic (hence unrealistic) conversations instead ;

because of the internal self-attention mechanism, they do not scale well to long conversations.

Therefore, a few recent works took a step back to borrow the best of both worlds (BoBW, i.e. multi-stage and end-to-end approaches) . The general principle is based on three steps:

split long conversations into shorter chunks;

apply end-to-end speaker diarization on each of them;

stitch them back together using speaker embeddings and unsupervised clustering.

To both start with a strong baseline and ensure reproducibility, we build on top of the speaker diarization pipeline available in version 2.1.1 of pyannote.audio open source toolkit that embraces this aforementioned three-steps principle. More precisely, we focus on the loss function used for training neural networks towards end-to-end speaker diarization of short audio chunks. calls this task speaker segmentation and states that working with short audio chunks has the following advantages:

it circumvents the scalability issue because the neural network only ever ingests short fixed-duration (5s in our case) audio chunks;

it makes it easier to train (as each conversation can now provide multiple training samples);

it allows to set a small upper bound on the (local) number of speakers $K_{\text{max}}$ – for instance, there is a 99% chance that any 5s chunk contains less than $K_{\text{max}}=3$ speakers in the 9 benchmarking datasets used in this paper.

Since its introduction a few years ago, the EEND line of work has been processing whole conversations with up to $K_{\text{max}}\approx 20$ (partially overlapping) speakers. Therefore, addressing it as a frame-wise multi-label classification problem was a brilliant (and efficient!) idea that we summarize in Section 2. Working with a much smaller number of ($K_{\text{max}}=3$) speakers opens up a new set of opportunities which have not yet been investigated by recent BoBW approaches. Inspired by , we propose to address speaker segmentation as a frame-wise powerset multi-class classification problem. Section 3 provides a detailed description of the approach but one can understand the main idea by looking at Figure 1. The advantages of this approach are two-fold: it completely removes the need for the (sensitive) detection threshold hyperparameter, and it leads to significant performance improvement of the overall speaker diarization pipeline (as shown by a thorough evaluation on 9 benchmarking datasets in Sections 4 and 5). In particular, we claim state-of-the-art performance on AISHELL-4 , AliMeeting , AMI , Ego4D , MSDWild , and REPERE benchmarks, make the approach available in the pyannote.audio library, and release pretrained models publicly.

From multi-label classification…

As depicted in the upper part of Figure 1, the reference segmentation of an audio chunk X can be encoded into a sequence of $K_{\text{max}}$-dimensional binary frames $\mathbf{y}=\{\mathbf{y_{1}},\ldots,\mathbf{y_{T}}\}$ where $\mathbf{y_{t}}\in\{0,1\}^{K_{\text{max}}}$ and $y_{t}^{k}=1$ if speaker $k$ is active at frame $t$ and $y_{t}^{k}=0$ otherwise. We may arbitrarily sort speakers by chronological order of their first activity but any permutation of the $K_{\text{max}}$ dimensions is a valid representation of the reference segmentation. Therefore, the binary cross entropy loss function $\mathcal{L}_{\text{BCE}}$ (classically used for such multi-label classification problems) has to be turned into a permutation-invariant loss function $\mathcal{L}$ by running over all possible permutations $\text{perm}(\mathbf{y})$ of $\mathbf{y}$ over its $K_{\text{max}}$ dimensions:

with $\mathbf{\hat{y}}=f(\mathbf{X})$ where $f$ is the segmentation model whose architecture is described below. In practice, for efficiency, we first compute the $K_{\text{max}}\times K_{\text{max}}$ binary cross entropy losses between all pairs of $\mathbf{y}$ and $\mathbf{\hat{y}}$ dimensions, and rely on the Hungarian algorithm to find the permutation that minimizes the overall binary cross entropy loss.

For the segmentation model $f$, we use the exact same architecture as the one introduced by . It consists in a SincNet convolutional block (that performs frame-wise feature extraction), four bi-directional LSTMs (that contains most of the learnable weights), two fully-connected layers, a classification layer with $K_{\text{max}}$ outputs, and a final sigmoid activation function to squash the frame-wise speaker activities $\hat{\mathbf{y}}$ between 0 and 1.

At test time, a subsequent binarization step is needed to decide whether each speaker $s$ is active in each frame $t$. This is achieved by comparing $\hat{\mathbf{y}}_{st}$ to a detection threshold $\theta\in\left[0,1\right]$ whose value needs to be tuned carefully.

… to powerset multi-class classification

While the standard multi-label loss relies on $K_{\text{max}}=3$ classes (one for each speaker $s_{1}$, $s_{2}$, $s_{3}$) that may be active simultaneously to encode overlapping speech, the powerset multi-class encoding considers a grand total of $K_{\text{powerset}}=7$ mutually exclusive classes:

$\{s_{1}\}$, $\{s_{2}\}$, $\{s_{3}\}$ for frames with one active speaker;

$\{s_{1},s_{2}\}$, $\{s_{1},s_{3}\}$, $\{s_{2},s_{3}\}$ for two overlapping speakersWe do not consider the case of three or more overlapping speakers because the number of frames for which this happens in our benchmarking datasets is marginal: $1.6\%$ in the compound dataset, $0.73\%$ in DIHARD III..

Switching to powerset encoding entails no radical change to the rest of the approach. It only amounts to changing the output size of the classification layer from $K_{\text{max}}=3$ to $K_{\text{powerset}}=7$, the final activation function from sigmoid to softmax, and the training loss function from binary to regular cross-entropy. Furthermore, a nice by-product is that the critical detection threshold $\theta$ is removed in favor of a simple argmax at test time.

Yet, powerset permutation space is slightly more complex than in the multi-label setting. For example, swapping speakers $s_{1}$ and $s_{2}$ in the multi-label case obviously implies swapping classes $\{s_{1}\}$ and $\{s_{2}\}$, but also overlapping speakers classes $\{s_{1},s_{3}\}$ and $\{s_{2},s_{3}\}$. Therefore, to find the optimal permutation in the powerset space, we use the following process:

convert target (y) from powerset to multi-label encoding;

convert binarized prediction ($\operatorname*{arg\,max}_{k}\hat{\textbf{y}}$) from powerset to multi-label encoding;

find the optimal permutation as described in Section 2;

permutate the target accordingly in the multi-label space, and convert it back to powerset encoding;

compute the cross-entropy loss in the powerset space.

While introduced the use of powerset encoding for speaker diarization, it has two main limitations. First, powerset encoding does not play well with a variable number of speakers: adding an extra speaker to a pool of $K$ speakers results in $K+1$ new powerset classes (one for the actual speaker, and one for each possible overlapping speakers classes). Therefore, had to resort to use a fixed and conservatively high number of speakers ($K_{\text{max}}=16$). Second, this leads to a number of powerset classes which is prohibitively high. For instance, using $K_{\text{max}}=16$ results in $K_{\text{powerset}}=\tbinom{0}{K}+\tbinom{1}{K}+\tbinom{2}{K}=137$ classes, most of which are almost never seen during training. For reasons already explained in the introduction, our approach overcomes both issues: we can stick to a small upper bound on the number of speakers ($K_{\text{max}}=3$) thanks to our processing of short audio chunks, and the unsupervised clustering step takes care of estimating the number of speakers specific to each conversation later in the pipeline.

Experiments

To ensure reproducibility, the research described in this paper builds on top of the speaker diarization pipeline from version 2.1.1 of pyannote.audio open source toolkit that follows the three-steps principle of BoBW approaches. More details on the pipeline internals are available in but we describe the gist of it in the following section.

The first step consists in applying the pretrained end-to-end neural speaker segmentation model pyannote/segmentation introduced in using a sliding window of 5s with a step of 500ms. A binarization step is further applied using the detection threshold $\theta\in\left[0,1\right]$ which constitutes the first hyperparameter of the approach. The second step consists in extracting one speaker embedding per active speaker in each 5s window. More precisely, speaker embeddings are only extracted from audio samples with exactly one active speaker: overlapping speech regions inferred automatically from the first step are discarded before computing the embeddings. All experiments reported in this paper rely on the pretrained ECAPA-TDNN model from SpeechBrain available at hf.co/speechbrain/spkrec-ecapa-voxceleb. The third and final step consists in applying agglomerative clustering on the aforementioned embeddings, using a second hyperparameter $\delta$ as the maximum allowed distance between centroids of two clusters for them to be merged. Once each local (i.e. from first step) active speaker is assigned to a global cluster, the global speaker diarization output can be constructed.

2 Datasets

We report performance on the official test sets of 9 different datasets: AISHELL-4 , AliMeeting , AMI (with two variants), DIHARD III , Ego4DAt the time of writing this paper, test labels of Ego4D were not available. Hence, Ego4D results were obtained on its development subset and should therefore be taken with a grain of salt. , MSDWild , REPERE , and VoxConverse . AMI variants are headset mix (or IHM, merged audio of the participants' headset microphones) and array 1 channel 1 (or SDM, first channel from the first far-field microphone array).

For datasets distributed with an official training/development split, we use the training subset as is. For those only providing a development set, we split it into two parts which we call training and development subsets in the rest of the paper. Models were trained on a compound multi-domain training set made of the concatenation of the training sets of all aforementioned datasets, but DIHARD. DIHARD was kept aside to measure how robust our models are to domain mismatch: it is composed of 11 sub-domains, all tested individually in this paper. Figure 2 showcases the amount of data from each dataset in the compound training set as well as their source and language. We do expect that including DIHARD in the compound training set would significantly improve the results found in this paper.

While Figure 2 highlights the relative imbalance between each dataset in the compound training set, we do make sure that the compound development set is balanced, in two variants:

dev_duration is built by picking random 5s chunks from the development subset of each dataset, for a grand total of 8 hours (one hour per dataset);

dev_files is built by randomly picking 25 files from the development subset of each datasetExcept for AISHELL-4, AliMeeting, and AMI which have only 20, 8, and 18 files in their respective development subsets., for a grand total of 164 files.

3 Experimental protocol

Speaker segmentation models are trained on the compound training set for at most one hundred hours, using Adam optimizer, an initial learning rate of $10^{-3}$, and a scheduler that divides the learning rate by $2$ after 30 epochs with no improvement. The validation metric used by the scheduler and for model selection is the local (i.e. computed on 5s chunks) diarization error rate on the dev_duration development set.

hyperparameters tuning is performed with pyannote.pipeline in order to minimize the average diarization error rate on the dev_files compound development subset. More specifically, in the case of an internal multi-label speaker segmentation model, two hyperparameters need to be optimized: the detection threshold $\theta$ and the clustering threshold $\delta$. In the case of powerset segmentation models, the detection threshold $\theta$ is not needed (because it is replaced by a parameter-less argmax): only the clustering threshold $\delta$ needs to be optimized.

Domain adaptation experiments reported in Section 5 are done by fine-tuning speaker segmentation models further on one specific domain. In those cases, for each dataset (or domain in case of DIHARD), model fine-tuning relies on the corresponding training set, model selection relies on the whole development set, and so does the tuning of pipeline hyperparameters.

Results and discussions

While the first line of Table 1 reports the performance of the speaker diarization pipeline based on the pretrained multi-label model available at hf.co/pyannote/segmentation, the second line is our (successful) attempt at retraining it from scratch with the proposed compound training set.

Powerset compound training consistently gets better results than its multi-label counterpart – with an average relative improvement of 8% for in domain datasets (from 25.6% to 23.5% diarization error rate). Further adapting the powerset model to each dataset gives an additional 8% relative improvement, reaching state of the art performance for most of them.

Both approaches are tested for robustness to domain mismatch using the DIHARD III dataset, which is composed of 11 domains that were never seen during (compound) training. Powerset training shows a $11\%$ relative improvement over multi-label (from 33.8% to 29.9%), even larger than for the in domain case. This suggests that powerset training is more robust to domain mismatch than its multi-label variant – a behavior that could be explained by the removal of the otherwise sensitive detection threshold $\theta$.

Figure 3 provides a detailed analysis of the errors committed by the various approaches. We find that most of the improvement comes from a consistent reduction in missed overlapping speech. This also holds for in-domain data (with missed detection rate reduced from 13.1% to 9.9% on average), supporting the idea that powerset explicit modeling of overlapping speakers classes has a huge impact on the overall performance.

Fine-tuning the powerset segmentation model on each DIHARD domain (`with domain adaptation' bars in Figure 3) tends to converge towards an overall better compromise between false alarm and missed detection rates, with little to no impact on speaker confusion. This suggests that this additional domain adaptation step also plays the role of the detection threshold $\theta$ (which is a separate hyperparameter in the multi-label case); thus making the whole approach even more robust.

Conclusion

In this paper, we study the impact of framing the speaker diarization task as a powerset multi-class classification problem. We extensively test this approach on 8 in domain datasets, as well as the 11 domains of DIHARD used as out-of-domain data. Compared to the classic multi-label approach, using a powerset multi-class encoding results in significant diarization error rate improvement (mostly due to better predictions on overlapping speech) and better robustness to domain mismatch. We obtain state of the art performance on AISHELL-4, AliMeeting, AMI, Ego4D, MSDWild, and REPERE. In the spirit of reproducible research, the powerset multi-class segmentation code is available in the open-source pyannote.audio library. The models trained on the compound training dataset, as well as their precomputed output on each separate dataset are available at github.com/FrenchKrab/IS2023-powerset-diarization.

References