Self-supervised Neural Factor Analysis for Disentangling Utterance-level Speech Representations

Consider an acoustic sequence ${\bf X}$ of $T$ frames. We denote ${\cal M}\subset\{1,\ldots,T\}$ as the index set indicating the frames in ${\bf X}$ to be masked. Define $\tilde{\bm{\mathbf{X}}}=\text{mask}(\bm{\mathbf{X}},{\cal M})$ as the masked version of $\bm{\mathbf{X}}$, where the masked $\bm{\mathbf{x}}_{t}$ $(t\in{\cal M})$ is replaced by a mask embedding. The BERT encoder $f_{\bm{\mathbf{\theta}}}(.)$ takes as input the masked sequence $\tilde{\bm{\mathbf{X}}}$ and outputs a feature sequence $\bm{\mathbf{H}}=[\bm{\mathbf{h}}_{1},\ldots,\bm{\mathbf{h}}_{T}]$. Let us introduce a $K$-dimensional binary random variable $\mathbf{y}_{t}$ for frame $t$ having a 1-of- $K$ representation, where $y_{tk}\in 0,1$ and $\sum_{k}y_{tk}=1$. Denote the output of the predictor as $q_{\phi}\left(y_{tk}\mid\mathbf{H}\right)$ . Given the target distribution for the masked frames $p\left(y_{tk}\right)$, the cross-entropy can be computed as:

However, we do not have access to the target distribution $p\left(y_{tk}\right)$. HuBERT solves this problem by iterative clustering to obtain the frame label $z_{tk}$ as a surrogate for $p\left(y_{tk}\right)$, where $z_{tk}\in 0,1$ and $\sum_{k}z_{tk}=1$. With the frame label $z_{tk}$, the cross-entropy loss can be re-written as:

At first, the cluster assignments are obtained by running K-means clustering on MFCCs. Then the model is updated by minimizing the masked prediction loss. New cluster assignments are obtained by running K-means on the updated features at the Transformer layer. The learning process then proceeds with new cluster assignments $\{\bm{\mathbf{z}}_{t}\}$. The masked prediction and cluster refinement are performed iteratively. The blue area in Figure 2 illustrates HuBERT’s masked prediction training.

Figure 1 shows that the K-means alignments can reveal meaningful speaker information. One simple way to obtain the utterance-level representation is to average the aligned frames in each cluster and concatenate the results. The probabilistic model for such approach can be written as follows:

where $\mathbf{h}_{t}^{i}$ is the Transformer layer features from the utterance $i$, $z_{tk}^{i}\in\{0,1\}$ is the frame label assigned by K-means, $\bm{\mu}_{k}$ is the $k$-th cluster center, $\bm{\mathbf{\Sigma}}_{k}$ is the covariance matrix of the $k$-th cluster, and $\bm{\mathbf{w}}_{k}^{i}$ is the utterance identity in the $k$-th cluster. The concatenation of $\bm{\mathbf{w}}_{k}^{i}$, i.e. $[\bm{\mathbf{w}}_{1}^{i},\ldots\bm{\mathbf{w}}_{K}^{i}]$, can be used as utterance identity representation. However, its dimension scales linearly with $K$. Instead, we decompose $\bm{\mathbf{w}}_{k}^{i}$ into the product of a cluster-dependent loading matrix $\bm{\mathbf{T}}_{k}$ and utterance identity vector $\bm{\mathbf{\omega}}^{i}$ for more compact representation:

Specifically, we train a K-means model using the Transformer layer features to produce $\{\bm{\mathbf{\mu}}_{k}\}$, which can be viewed as content representations of the speech. Then, we run K-means to produce frame labels $\{z_{tk}^{i}\}$ and calculate $\{\bm{\mathbf{\Sigma}}_{k}\}$ and cluster weight prior $\{\pi_{k}\}$ for the $K$ clusters, which we denoted as $\bm{\mathbf{\Phi}}=\{\pi_{k},\bm{\mathbf{\mu}}_{k},\bm{\mathbf{\Sigma}}_{k}\ |k=1,\ldots,K\}$. With cluster parameters and frame labels $\{z_{tk}^{i}\}$, we only have one set of parameters $\{\bm{\mathbf{T}}_{k}\}$ and one latent variable $\bm{\mathbf{\omega}}^{i}$ left in the model, which is a problem that can be solved with the expectation-maximization (EM) algorithm.

Given a sequence of frame-level features $\bm{\mathbf{H}}^{i}=\{\bm{\mathbf{h}}_{1}^{i},\ldots,\bm{\mathbf{h}}_{T}^{i}\}$, the frames labels (alignments) $\bm{\mathbf{Z}}^{i}=\{z^{i}_{tk}|t=1,\ldots,T;k=1,\ldots,K\}$, and cluster parameters $\bm{\mathbf{\Phi}}$, we can use the EM algorithm to find $\bm{\mathbf{T}}=\{\bm{\mathbf{T}}_{k}|k=1,\ldots,K\}$. In the E-step, we compute the posterior of utterance identity $\bm{\mathbf{\omega}}^{i}$:

where $\bm{\mathbf{z}}^{i}_{t\bullet}=\{{z}^{i}_{tk}\}_{k=1}^{K}$ and $p_{\bm{\mathbf{T}}}\left(\bm{\mathbf{\omega}}^{i}|\bm{\mathbf{H}}^{i};\bm{\mathbf{Z}}^{i},\bm{\mathbf{\Phi}}\right)$ is the probability distribution of $\bm{\mathbf{\omega}}^{i}$ conditioned on $\bm{\mathbf{H}}^{i}$ given $\bm{\mathbf{Z}}^{i}$ and $\bm{\mathbf{\Phi}}$. Because the alignments $\bm{\mathbf{Z}}^{i}$ and the cluster parameters $\bm{\mathbf{\Phi}}$ are fixed while optimizing the likelihood, we drop the dependency when expressing the posterior for simplicity.

In the M-step, we choose the $\bm{\mathbf{T}}$ that maximize the expected log-likelihood:

where $\bm{\mathbf{T}}^{{}^{\prime}}$ is the loading matrix from the previous M-step (or randomly initialized). Eq. 6 has a closed-form solution. After the matrix $\bm{\mathbf{T}}$ is found, the mean of the posterior $\mathop{\mathbb{E}}[\bm{\mathbf{\omega}}|\bm{\mathbf{H}}]$ is used as the utterance identity representation.

Learning via gradient on ELBO There are two limitations to learning matrix $\bm{\mathbf{T}}$ using the EM algorithm. First, the EM algorithm limits the possibility of large-scale training. In Eq. 6, the loading matrix $\bm{\mathbf{T}}$ is estimated using the whole training set, contrary to the stochastic update in modern DNN training. Another disadvantage is the separation between the Transformer layers and the FA model during training, which prevents the possibility of joint optimization of the matrix $\bm{\mathbf{T}}$ and Transformer layers’ parameters $\bm{\mathbf{\theta}}$.

We aim to derive a learning rule that is amenable to stochastic updates and allows joint optimization of the FA model and the Transformer layers. As a latent variable model, the log-likelihood of our FA model can be written as (Bishop & Nasrabadi, 2006; Kingma & Welling, 2013):

where $\mathcal{L}_{\text{ELBO}}\left(\bm{\mathbf{H}}^{i};\bm{\mathbf{T}}\right)$ is called the evidence lower bound (ELBO). $D_{\text{KL}}\left(q(\bm{\mathbf{\omega}}^{i})\|p_{\bm{\mathbf{T}}}(\bm{\mathbf{\omega}}^{i}|\bm{\mathbf{H}}^{i})\right)$ is the KL-divergence between the approximate posterior $q(\bm{\mathbf{\omega}}^{i})$ and true posterior $p_{\bm{\mathbf{T}}}(\bm{\mathbf{\omega}}^{i}|\bm{\mathbf{H}}^{i})$. Minimizing KL or maximizing the ELBO can both increase the log-likelihood. In the case of our model, minimizing the KL is easy as the posterior of $\bm{\mathbf{\omega}}$ is tractable, which gives rise to the E-step in Eq. 5. To optimize the ELBO, we need to re-write Eq. 8 as:

Because we already know the closest ELBO to likelihood is when $q(\bm{\mathbf{\omega}}^{i})$ equals to the posterior $p_{\bm{\mathbf{T}}}\left(\bm{\omega}^{i}\mid\mathbf{H}^{i}\right)$, Eq. 9 can be written as:

where $\bm{\mathbf{T}}^{{}^{\prime}}$ is the loading matrix from the last update. We can see the first term is a constant with respect to $\bm{\mathbf{T}}$. Therefore, the gradient of the lower-bound with respect to $\bm{\mathbf{T}}$ is:

The gradient with respect to the Transformer features $\frac{d\mathcal{L}_{\text{ELBO}}}{d\bm{\mathbf{H}}^{i}}$ involves both terms in Eq. 10:

By applying the chain rule, we can obtain the gradient with respect to the Transformer parameters $\bm{\mathbf{\theta}}$:

Eq. 13 shows that we can backpropagate the gradient of ELBO back to the Transformer layers. The total loss of our NFA model is:

Therefore, in addition to HuBERT’s mask prediction and self-training, in each forward pass, we will compute the posteriors $p_{\bm{\mathbf{T}}}\left(\bm{\omega}^{i}\mid\mathbf{H}^{i}\right)$ (Eq. 5) given a sequence of BERT features and frame labels produced by K-means. Then, we use the posteriors to evaluate the gradient with respect to $\bm{\mathbf{T}}$ to update the loading matrix and the gradient with respect to BERT features $\bm{\mathbf{H}}^{i}$ to update the SSL model parameters $\bm{\mathbf{\theta}}$. Algorithm 1 summarizes the whole training procedure of our NFA.

Speech Tasks and Datasets The speech tasks that we will evaluate include:

• •

  Automatic speaker verification (ASV or SV), speaker identification (SID), and speaker diarization (SD). We followed the SUPERB protocol (Yang et al., 2021) using the VoxCeleb1 (Nagrani et al., 2017) training split to train the model and used the test split to evaluate speaker verification performance. Note that the reported ASV downstream model in (Yang et al., 2021) is a deep neural network (Snyder et al., 2018) trained on SSL features (Yang et al., 2021). The evaluation metric is equal error rate (EER) (the lower, the better). For speaker identification, we used the VoxCeleb1 train-test split provided by the SUPERB organizer. The evaluation metric is accuracy. For SID, the SUPERB downstream model is a linear classifier trained on averaged SSL features. Speaker diarization is to segment and label a recording according to speakers. We followed the SUPERB protocol using the LibriSpeech (Panayotov et al., 2015) splits for training and evaluation. The SUPERB downstream model is a recurrent neural network. The evaluation metric is diarization error rate (DER) (the lower, the better)

• •

  Emotion recognition (ER). We used IEMOCAP (Busso et al., 2008) dataset. Following the same protocol as SUPERB, we dropped the unbalance emotion classes to leave the neutral, happy, sad, and angry classes. The evaluation metric is accuracy. The SUPERB downstream model is a linear classifier trained on averaged SSL features.

• •

  Language identification (LID). Language identification is not included in the SUPERB benchmark. We included it because it is also an important utterance-level task. The dataset we used is the the Common Language dataset prepared by (Sinisetty et al., 2021), which includes 45 languages with 45.1 hours of recordings. On average, each language has one-hour recordings.¹¹1https://huggingface.co/datasets/common_language The downstream baseline is a linear classifier trained on averaged SSL features.

Pre-trained models The pre-trained models we used in this paper include HuBERT (Hsu et al., 2021a), WavLM (Chen et al., 2022), and wav2vec2-XLS-R (Babu et al., 2022). HuBERT and WavLM models were used in speaker and emotion evaluation. Because language identification requires models trained on multi-lingual data, wav2vec2-XLS-R was used.

Implementation details. The HuBERT and Wav2vec2-based NFA models were trained on LibriSpeech using the model checkpoints provided by fairseq. The language identification NFA models were trained on the Common Language dataset using the XLS-R checkpoint. $\lambda$ in Eq. 14 is set to 0.01 for all models. After the optimization steps in Algorithm 1 were done, we re-trained the loading matrix $\bm{\mathbf{T}}$ for each task with EM using unlabeled task-related data. Other than specifically stated, the acoustic features were extracted from layer 6 for the base SSL models (HuBERT, WavLM, and Wav2Vec2-XLS-R) and layer 9 for the large SSL models. The number of clusters in K-means is 100, and the rank of loading matrix dimension is 300 for all NFA models. After utterance-level representations have been extracted using Eq. 7, we used the simple logistic classifier in sklearn (Pedregosa et al., 2011) for SID, ER, and LID. For speaker verification, we used the PLDA backend. For SD, we used linear discriminant analysis (LDA) to reduce the dimension to 200 and then used agglomerative hierarchical clustering to produce speaker assignments. Note that all our downstream methods are linear models.

In this section, we evaluate the NFA’s performance on SUPERB tasks (Yang et al., 2021; Chen et al., 2022). Besides the standard speaker-related and emotion recognition, we also included language identification (LID) on Common Langue (Sinisetty et al., 2021). For LID, we followed the same protocol as other SUPERB tasks, i.e., the SSL models’ weights were frozen, and only linear models were trained with labeled data without data augmentation. To give a better idea of the expected performance of each task in unrestricted settings, we also included the results using the fine-tuned SSL models on the ASV and ER tasks and the current best result in the Common Language dataset reported by other researchers.

The results are presented in Table 1. As observed in the table, NFA significantly outperforms all SSL models across ASV, SD, SID, and LID. NFA performs only marginally worse than the self-supervised Conformer (Shor et al., 2020), which has been specifically designed for utterance-level tasks. In speaker verification, the relative EER reduction is 40% when compared with the WavLM, the previous best model on utterance-level tasks. It is worth noting that WavLM’s ASV baseline used a DNN network trained on the Transformer features, but we only use linear models. Our models even perform better than the fully fine-tuned models in (Wang et al., 2021) in both ASV and ER tasks. For LID, our XLS-R-based NFA performs better than the best-reported result on Common Language by SpeechBrain (Ravanelli et al., 2021).

One of the most attractive features of wav2Vec and HuBERT is their performance on low label-resource ASR. The resource efficiency of these models enables the potential development of many low label-resource languages and speech tasks where labeled data are hard to collect. In this section, we evaluate NFA performance in low label-resource settings. To this end, we divided the labeled dataset in the speaker recognition, emotion recognition, and language identification tasks into 10%, 20%, and 30% subsets as low label-resource settings. For ASV, SD, SID, and ER, we extracted the embeddings from a large Hubert-based NFA model. For LID, we used the embeddings from the XLS-R-based NFA model. WavLM Large and XLS-R were used as performance references. To reduce the performance deviation in the division, we ran each partition five times and reported the results. The loading matrices in the NFA models were trained using the entire unlabeled dataset. The results are presented in Figure 3.

We can see that even with only 10% of labeled data for the downstream models, NFA’s performance in ER, SID, and LID is very close to the WavLM and XLS-R. For ASV and SD, our method already outperforms the WavLM models trained on fully labeled data. With 20% labeled data, NFA already outperforms WavLM and XLS-R on all tasks. This shows the high resource efficiency of our NFA models.

In Figure 1, we observe that by clustering and aligning the Transformer features, speaker information can be revealed. This is all done without labeled data. But how discriminative these unsupervised learned embeddings are? We will evaluate NFA embeddings’ zero-shot performance quantitatively in this section. Specifically, we evaluated NFA models on zero-shot speaker verification. After we extracted the utterance-level representations using Eq. 7, we directly used cosine similarity to obtain verification scores without any supervised training (the models were never given speaker information). We evaluated the performance on (1) LibriSpeech, which is considered in-domain data as HuBERT and NFA were trained on this dataset (Panayotov et al., 2015; Hsu et al., 2021a), (2) Voxceleb1-test, a popular speaker verification dataset (Nagrani et al., 2017), and (3) VOiCES (Nandwana et al., 2019), a dataset used to evaluated speaker verification robustness against noise and room reverberation. As a comparison, we also included i-vector (Dehak et al., 2010) and averaged Transformer features (HuBERT rows in Table 2) as baselines.

The results are presented in Table 2. Without supervision, simple averaging the Transformer features cannot produce useful speaker representations. It even performs worse than i-vector, a non-DNN approach. NFA embeddings, however, achieve an EER of 3.98% on LibriSpeech without any supervised training. This suggests that during self-supervised learning, the model has already learned to differentiate speakers, which also empirically demonstrates that the NFA model can disentangle speaker information from the content information. However, when evaluated on VoxCeleb1 and VOiCES, the performance of zero-shot SV dropped significantly. This may be because VoxCeleb1 and VOiCES are real-world speech datasets containing spontaneous speech and environmental noise. NFA and HuBERT were pre-trained on a read speech dataset. The domain discrepancy in SSL models can have a significant impact on the downstream tasks, as mentioned in (Hsu et al., 2021b). Another interesting observation is that scaling the model size improves the zero-shot SV performance, as shown when using HuBERT Large and NFA large models.

Because our NFA models show excellent zero-shot performance, we can use them to evaluate the discrimination power from each Transformer layer before supervised learning is applied. We extracted the acoustic features from Layer 1 to Layer 12 of the Transformer in the NFA model to conduct zero-shot speaker verification and language identification. For language identification, we used top-1 accuracy as the metric. Then, we used the labeled data to train an LDA on top of NFA embeddings to compare the results. The results are presented in Figure 4.

The blue lines in Figure 4 show that under zero-shot settings, both speaker and language discriminative abilities increase from Layer 1 up to Layer 6. Then, the features from the deeper layers have poorer performance. This is largely consistent with the supervised baselines (orange lines), with Layer 7 obtaining the lowest speaker verification error and Layer 6 having the highest language identification top-1 accuracy in supervised settings. This shows that our NFA models’ zero-shot performance can be a reliable predictor of supervised performance.

To assess whether gradient-based learning has an edge over the Expectation-Maximization (EM) method, we extracted HuBERT features and separately trained a factor analysis model using EM. The results are displayed in Table 3. We observe that gradient-based optimization consistently outperforms EM-based I-vector trained on HuBERT features. This suggests that jointly training the NFA model with the SSL model can yield more potent feature representations than training the two modules independently.

The ultimate goal of a self-supervised learning (SSL) speech model is to utilize a single backbone model for all downstream tasks. Consequently, it’s critical that the NFA model does not compromise performance on content-based tasks such as ASR. To ensure this, we compared the performance of the NFA and the large NFA model against HuBERT on the LibriSpeech clean subset. The results, as shown in Table 4, demonstrate that the NFA and large NFA models perform on par with HuBERT. This confirms that our NFA model does not sacrifice performance on content-based tasks.