Speech Separation based on pre-trained model and Deep Modularization

Several unsupervised speech separation methods have been developed in recent years. One notable technique is MixIT (Wisdom et al., 2020b), which performs speech separation by mixing mixtures. Given a set of mixed speech signals $X={x_{1},x_{2},\cdots,x_{n}}$, where each mixture $x_{i}$ contains up to $N$ sources, MixIT creates a mixture of mixtures (MoM) by randomly drawing and adding mixtures from $X$. For instance, if $x_{1}$ and $x_{2}$ are selected, the MoM $\bar{x}$ is created as $\bar{x}=x_{1}+x_{2}$. This MoM is then fed into a deep neural network (DNN), which estimates the sources $\hat{s}$ within $x_{1}$ and $x_{2}$. The model is trained to minimize the following loss:

Here, $A\in B^{2\times M}$ is a set of binary matrices, where each column sums to 1, and $M$ is the maximum number of sources in a given mixture $x_{i}$. The loss minimizes the difference between mixtures $x_{i}$ and the remixed separated sources $A\hat{s}$. A limitation of MixIT is over-separation, where the model estimates more sources than are present in the mixtures (Zhang et al., 2021). Additionally, MixIT struggles with denoising (Saito et al., 2021). Variants of MixIT, such as (Zhang et al., 2021) and (Karamatlı & Kırbız, 2022), address the over-separation problem, while (Trinh & Braun, 2022) and (Saito et al., 2021) enhance MixIT for denoising tasks.

Another unsupervised method, RemixIT (Tzinis et al., 2022), employs a teacher-student model for speech denoising and separation. The teacher model estimates the clean speech sources $\hat{s}_{i}$ and noises $\hat{n}_{i}$ for a batch of size $b$. The estimated noises are randomly mixed to generate $n^{p}$, which, together with the clean sources, form new mixtures $\hat{m}_{i}=\hat{s}_{i}+n^{p}$. These new mixtures are used as input to the student model, which is trained to output $\hat{s}_{i}$ and $n^{p}$, effectively denoising the speech. RemixIT relies on a pre-trained speech enhancement model as the teacher.

Other unsupervised approaches, such as (Wang et al., 2016), perform speech separation based on speaker characteristics like gender. These methods use i-vectors to model differences in vocal tract properties, pitch, timing, and rhythm between male and female speakers, and the DNN acts as a gender separator.

Speech separation tools are designed to accept either Fourier spectrum or time-domain input features. Fourier spectrum-based models apply the discrete Fourier transform (DFT) to convert the raw time-domain signal into the frequency domain. These models recognize the non-stationary nature of speech and its variability in both time and frequency. Commonly used features include log-power spectrum (Fu et al., 2017; Xu et al., 2015), Mel-frequency spectrum (Liu et al., 2022; Du et al., 2020; Donahue et al., 2018), DFT magnitude (Nossier et al., 2020; Grais & Plumbley, 2018; Fu et al., 2019), and complex DFT features (Fu et al., 2017; Williamson & Wang, 2017; Kothapally & Hansen, 2022a; b). Most models using Fourier spectrum features assume that phase information is less critical for human auditory perception and primarily focus on magnitude or power for training (Xu et al., 2014; Kumar & Florencio, 2016; Du & Huo, 2008; Tu et al., 2014; Li et al., 2017). The noisy phase is typically used for signal reconstruction based on the work of (Ephraim & Malah, 1984), which suggests that the phase of the noisy signal is a close estimator of the clean signal’s phase.

However, recent research (Paliwal et al., 2011) shows that further improvements in speech quality can be achieved by processing both the magnitude and phase spectra. To address phase challenges with Fourier-based features, several time-domain speech separation models have been developed, such as (Luo & Mesgarani, 2018; Luo et al., 2020; Luo & Mesgarani, 2019a; Subakan et al., 2021a). These models replace the DFT-based input with data-driven representations learned during model training. Instead of converting the raw signal to the frequency domain, they process the mixed waveform directly, generating either the estimated clean sources or masks applied to the noisy waveform to produce clean speech. This approach allows the models to fully capture both magnitude and phase information, leading to improved separation performance (Luo et al., 2020).

To extract meaningful representations from input speech, we employ a pre-trained model capable of processing either T-F bin patches or time-domain blocks. The model is trained on a large corpus of patches or speech blocks, enabling it to generate robust feature representations. These learned features serve as input to the subsequent clustering step.

We leverage contrastive learning (CL) to train the pre-trained model $f_{\theta}$, where the goal is to minimize the distance between similar objects (positive pairs) and maximize the distance between dissimilar objects (negative pairs) in the embedding space. Specifically, CL defines a representation function $f_{\theta}:x\mapsto\mathbb{R}^{d}$, mapping augmented inputs into $d$-dimensional vectors. The objective is to ensure that similar augmentations of the same input are close in the embedding space, while different inputs are pushed apart.

Typically, augmentations $(x,x^{+})$ are generated by applying two distinct augmentation functions to the same input. In our case, these augmentations preserve critical features, such as speaker identity, while altering irrelevant aspects like noise and reverberation.

The contrastive learning model $f_{\theta}$, parameterized by $\theta$, is trained using either T-F bin patches or raw speech blocks as inputs. Positive pairs consist of two T-F bin patches or raw blocks from the same speaker, while negative samples are patches or blocks from different speakers. For example, a speaker similarity function $\mathcal{T}$ randomly selects a positive pair $(x_{n_{i}},x_{nr_{i}})$, where $x_{n_{i}}\in\bar{X}_{n}$ and $x_{nr_{i}}\in\bar{X}_{nr}$, ensuring both come from the same speaker.

For a batch of size $N$, given a positive pair $(x_{n_{i}},x_{nr_{i}})$, the remaining $N-2$ patches or blocks in the batch serve as negative samples. The contrastive learning objective is optimized via the following loss function:

Here, $\text{sim}(z_{n_{i}},z_{nr_{i}})$ represents the similarity between the embeddings of the positive pair, and $\tau$ is a temperature parameter controlling the separation between positive and negative pairs.

Algorithm 1 outlines the training process of the pre-trained model using contrastive learning.

The goal is to use a deep modularization technique to cluster patches or raw speech blocks such that those dominated by a given speaker are clustered together. Deep modularization is a technique that integrates a graph clustering objective with a deep neural network (DNN). We start by defining a graph $G(V,E)$ where $V=(v_{1},v_{2},\cdots,v_{n})$, with $|V|=n$, is the set of all patches or raw speech blocks generated from a mixed speech signal, as described in the speech pre-processing section. The set of edges is $E\subset V\times V$, with $|E|=m$, which connects the generated patches or raw speech blocks (subsequently referred to as "blocks"). The adjacency matrix of $G$ is denoted as $A$, where $A_{ij}=1$ if ${v_{i},v_{j}}\in E$, and 0 otherwise. The degree of $v_{i}$ is defined as $d_{i}=\sum_{j}^{n}A_{ij}$. Our objective is to generate a graph partition function $\mathcal{F}:V\rightarrow{1,\cdots,k}$ that splits the set of patches or blocks $V$ into $k$ partitions, where $v_{i}={v_{j}:\mathcal{F}(v_{j})=i}$, given the patches or block attributes $\bar{F}\in R^{n\times d}$ learned from the pre-trained model.

To partition the vertices, we use the statistical approach known as modularity ($Q$) (Newman, 2006). Modularity measures the difference between the actual number of edges within partitions and the expected number of edges in equivalent randomized partitions (null networks). Formally, modularity is defined as:

A high value of $Q$ indicates stronger similarity among members of the same partition. Therefore, the goal is to maximize $Q$. Let $g_{i}$ represent the community to which vertex $i$ belongs. Modularity ($Q$) is derived in (Newman, 2006) as:

where $\delta(g_{i},g_{j})$ equals 1 if vertices $i$ and $j$ belong to the same partition and 0 otherwise. $P_{ij}$ is the expected number of edges between vertices $i$ and $j$, while $A_{ij}$ is the actual number of edges between them. If vertices $i$ and $j$ have degrees $d_{i}$ and $d_{j}$, respectively, then the expected degree of vertex $i$ is $\sum_{j}P_{ij}=d_{i}$. Hence, vertex $i$ and vertex $j$ are connected with probability $P_{ij}=\frac{d_{i}d_{j}}{2m}$ (Newman, 2006). Consequently, Equation 2 becomes:

Maximizing $Q$ is NP-hard (Brandes et al., 2006). To solve this, we define a partition assignment matrix $S\in R^{n\times k}$, where $n$ is the number of vertices. Each column of $S$ indexes a partition, meaning $S={s_{1}\mid s_{2}\mid\cdots\mid s_{k}}$, with $S_{ij}=1$ if vertex $i$ belongs to partition $j$, and 0 otherwise. The columns of $S$ are mutually orthogonal, as each row sums to 1, and $S$ satisfies the normalization $Tr(S^{T}S)=n$, where $Tr(\cdot)$ is the matrix trace. Based on this definition, $\delta(g_{i},g_{j})=\sum_{k=1}^{k}S_{ik}S_{jk}$, and Equation 3 becomes:

where $B$ is the modularity matrix, defined as $B_{ij}=A_{ij}-P_{ij}$. By relaxing $S\in R^{n\times k}$, the optimal $S$ that maximizes $Q$ corresponds to the top $k$ eigenvectors of matrix $B$. Work by (Müller, 2023) enhances this objective by adding a regularization term to avoid trivial partitions, as shown in Equation 5. This regularization maintains consistency in community detection as the number of nodes increases (Müller, 2023):

Here, $|\cdot|_{F}$ is the Frobenius norm. The complexity of the modularity term $Tr(S^{T}BS)$ is $\mathcal{O}(n^{2})$ per update of $L(S)$, making the training computationally expensive. To address this, (Müller, 2023) proposes to decompose $Tr(S^{T}BS)$ into sparse matrix-matrix multiplication and rank-one degree normalization, reducing the complexity to $\mathcal{O}(d^{2}n)$:

We adopt the objective in Equation 6 to modularize the features $\bar{F}\in R^{n\times d}$ learned via the pre-trained model. To optimize the cluster assignment, we adapt the deep neural network graph partitioning technique proposed in (Bianchi et al., 2020) and (Müller, 2023), where nodes of a graph are partitioned as follows:

Here, $\tilde{A}=D^{-\frac{1}{2}}AD^{-\frac{1}{2}}$, $X$ are the input features, $D$ is the diagonal matrix of degrees $d_{1},\dots,d_{n}$, and $A$ is the adjacency matrix. In Equation 7, node features $\bar{F}$ are learned via a graph neural network (GNN), and the assignment matrix $S$ is generated through a softmax function. In (Bianchi et al., 2020), the assignment matrix $S$ is established via a multilayer perceptron (MLP) with softmax at the output layer. In our case, we formulate the problem as:

Here, the feature matrix $\bar{F}$ is obtained via the pre-trained model (Pre-M). Since we extracted $3\times 3$ patches and speech separation is typically performed over individual $1\times 1$ T-F bins via masking, we apply a Gaussian-weighted patch averaging technique to compute the mean feature representation for each $1\times 1$ T-F bin. To preserve local dependencies, we restrict the averaging to the three closest neighboring T-F bins, i.e., $|C|=3$, in both the time and frequency domains. For each feature vector $\bar{f}_{j}\in\bar{F}$, the Gaussian-weighted averaged feature is computed as:

where $C$ consists of the central T-F bin and its two closest neighbors. The weights $w_{i}$ are assigned using a 2D Gaussian kernel, emphasizing the importance of proximity within the local T-F region to retain crucial local information.

The averaged feature $\bar{f}_{\text{avg},j}$ provides a compact and locally representative feature that captures the most relevant information from the neighboring T-F bins and serves as input to the BLSTM model.

The BLSTM, as proposed in Huang et al. (2022), takes the averaged patch feature and assigns it to the $j$-th column of the cluster assignment matrix $S$. The final cluster assignments are established using a softmax at the output layer.

To optimize the assignment $S$, we use the loss function described in Equation 6, effectively training a DNN model using a graph clustering objective but with averaged patch features. This enhances the robustness of the learned features by reducing noise and capturing more meaningful representations for each patch.

For speech blocks, no further averaging or processing is applied, and the features $\bar{F}$ are used directly for clustering or downstream tasks.

To construct the adjacency matrix $A$, we first measure the similarity between each T-F bin or block $i$ and all other nodes in the feature space. This similarity is computed using the inner product:

where $i,j=1,2,\dots,n$, and $\bar{f}_{i}$ and $\bar{f}_{j}\in\bar{F}$ are feature vectors obtained from the pre-trained model (Pre-M). The higher the inner product, the more similar the two feature vectors are.

Next, to form the edges of the graph, we select a similarity threshold $\theta$. If $e_{ij}<\theta$, we remove the edge between nodes $i$ and $j$, thereby limiting the graph to meaningful edges. The adjacency matrix is defined as:

Choosing an optimal threshold $\theta$ is crucial for constructing a meaningful graph. A high value of $\theta$ ensures only highly similar nodes are connected, but this can exclude meaningful relationships. On the other hand, a low $\theta$ results in a dense graph filled with weaker, uninformative edges, which increases computation time during clustering.

To find the optimal $\theta$, we experimented on multiple datasets, varying the threshold values and observing both modularity and the number of clusters generated. Modularity measures the quality of clusters, with higher values indicating better-defined clusters.

Figure 2 shows how modularity and the number of clusters change with varying $\theta$ on the WSJ0-5mix test dataset. For visualization purposes, modularity values are normalized (multiplied by 100) and the number of clusters is scaled by 10. As shown in the figure, increasing $\theta$ boosts modularity but risks creating singleton partitions. Decreasing $\theta$, however, increases unnecessary clusters and lowers modularity.

After careful evaluation, we selected $\theta=0.3$ as the optimal threshold, which was used consistently across all our experiments.

We first evaluated which of the four model configurations (MOD1, MOD2, MOD3, MOD4) produces better clusters (partitions). To measure cluster quality, we used the graph-based cluster evaluation metrics proposed in (Yang & Leskovec, 2012), focusing on metrics that quantify how well a partition is separated from the rest, i.e., the number of edges connecting a given partition to other partitions. A high-quality partition should have few edges pointing outside. The most relevant metrics for our study are graph modularity and conductance.

Cluster Conductance (C) is defined as $C=\frac{c_{s}}{2m_{s}+c_{s}}$, where $S$ is a partition, $m_{s}$ is the number of edges within $S$ (i.e., $m_{s}=\{(u,v)\in E\mid u\in S,v\in S\}$), and $c_{s}$ is the number of edges on the boundary of $S$ (i.e., $c_{s}=\{(u,v)\in E\mid u\in S,v\notin S\}$). Conductance measures the fraction of edges pointing outside a partition, with lower values indicating better quality.

Graph Modularity ($\mathcal{Q}$) is defined as $\mathcal{Q}=\frac{1}{4}(m_{s}-E(m_{s}))$, where $E(m_{s})$ is the expected value of $m_{s}$. Higher modularity indicates better partition quality.

The results of our evaluation are shown in Table 2. Comparing similarly trained models across all datasets, we found that encoders trained on raw speech blocks produced higher-quality clusters than those trained on T-F bin patches. This suggests that for this setup, raw speech features capture more relevant speech information compared to STFT-transformed features. Additionally, fine-tuning the encoder significantly improved cluster quality compared to using a frozen encoder.

We evaluated the performance of four speech separation models (MOD1, MOD2, MOD3, and MOD4) using the WSJ0-2mix, WSJ0-3mix, and WSJ0-4mix test datasets. The results, summarized in Table 3, show that MOD4 consistently achieved the highest performance across all metrics, followed by MOD2, irrespective of the number of speakers in the mixture. Specifically, MOD4 demonstrated significant improvement in SDRi (Signal-to-Distortion Ratio improvement), SI-SNRi (Scale-Invariant Signal-to-Noise Ratio improvement), and STOI (Short-Time Objective Intelligibility), outperforming the other models across all test sets.

Interestingly, this contrasts with the clustering quality evaluation, where MOD4 also produced the best clusters, but MOD3 outperformed MOD2 in cluster quality. The lower performance of MOD3 compared to MOD2 in separation quality can be attributed to phase reconstruction issues that arise when using STFT-transformed features. The phase reconstruction process, essential for high-fidelity speech reconstruction, is more challenging with MOD3’s approach, leading to suboptimal separation performance.

Across the three datasets, the superior performance of MOD4 and MOD2 is consistent with models trained on raw speech blocks for both pre-training and downstream tasks. This indicates that raw speech blocks capture more relevant and meaningful features for speech separation compared to models trained with STFT-transformed patches (T-F bins). The inferior performance of the T-F bin-based models (MOD1 and MOD3) likely stems from the phase reconstruction difficulties inherent in this approach, which is further exacerbated in the separation process.

We evaluated the performance of the proposed speech separation technique (MOD4) against several state-of-the-art models on the WSJ0-2mix and WSJ0-3mix test datasets. The results, summarized in Table 4, demonstrate that MOD4 consistently achieves competitive results. On the WSJ0-2mix dataset, MOD4 attained an SI-SNRi of 22.6, trailing MossFormer2 by 1.5 points, and an SDRi of 22.9, which is only 0.6 points lower than the top-performing TF-GridNet. These results indicate that MOD4, despite not relying on parallel datasets, is highly competitive with fully supervised methods.

On the WSJ0-3mix dataset, MOD4’s robustness and scalability to higher-source mixtures become evident. While MossFormer2’s SI-SNRi dropped by 1.9 points between the WSJ0-2mix and WSJ0-3mix tasks, MOD4 only experienced a minor drop of 0.8 points. This ability to maintain performance across more complex mixtures highlights MOD4’s generalization capability, with MOD4 achieving an SI-SNRi of 21.8, only 0.4 points behind MossFormer2, showing strong scalability as the number of sources increases.

Furthermore, MOD1 and MOD3 outperformed $\mathcal{L}{MISI-5}^{Enh3}$, a supervised STFT-based method, which uses non-negative masking with MISI for phase reconstruction. The superior performance of MOD1 and MOD3 suggests that the proposed models’ ability to more accurately predict the magnitude of the source signal results in higher-quality separation, even in the absence of parallel data. This highlights the strength of raw speech block models over STFT-based models for speech separation tasks.

We evaluated the performance of the proposed technique on mixtures with a large number of sources. The results, presented in Table 5, demonstrate that MOD4 achieves superior performance across the SDRi metric, outperforming existing tools by 0.2, 0.3, and 2.0 points on the WSJ0-5mix, Libri5Mix, and Libri10Mix datasets, respectively. This consistent performance highlights the ability of the proposed technique to scale effectively to high-source mixtures while maintaining high-quality separation.

Specifically, in the WSJ0-5mix test set, MOD4 achieved the highest SDRi score of 13.4, surpassing the previous top-performing Hungarian algorithm by 0.2 points. Similarly, in the Libri5Mix test set, MOD4 obtained an SDRi score of 13.0, outperforming Hungarian by 0.3 points. Notably, the most significant improvement was observed in the Libri10Mix test set, where MOD4 outperformed Hungarian by 2.0 points, further demonstrating its scalability in handling increasingly complex mixtures.

These results suggest that MOD4’s architecture and underlying techniques are well-suited for separating a larger number of sources, surpassing several state-of-the-art methods. This scalability is especially important for real-world applications, where mixtures of many sources are common, and maintaining high separation quality becomes increasingly challenging.

We also assessed the proposed technique on two datasets that include noise and reverberation: WHAM! and WHAMR!. The results, summarized in Table 6, show that MOD4 achieved the best performance compared to other state-of-the-art methods on both datasets. On the WHAM! dataset, which requires simultaneous speech separation and denoising, the model was able to effectively handle the additional noise partition, demonstrating its robustness in noisy environments. Similarly, on the WHAMR! dataset, where the model must handle denoising, dereverberation, and speech separation simultaneously, MOD4 outperformed competing models across both the SI-SNRi and SDRi metrics.

In this experiment, we replaced deep modularization with classical $k$-means clustering to generate partitions, setting $k=2$ for the WSJ0-2mix test dataset and $k=3$ for the WSJ0-3mix test dataset. Specifically, after extracting T-F bin patches or raw speech block features using a frozen pre-trained encoder, we applied $k$-means to cluster these features and generate partitions, which were subsequently used for mask generation. We evaluated $k$-means in two configurations:

1. 1.

   $k$-means (raw speech blocks): Clustering was applied to features extracted from raw speech blocks.

2. 2.

   $k$-means (patches): Clustering was applied to features extracted from patches.

The mask generation and speech reconstruction processes remained identical to those outlined in Section 3. We compared the performance of $k$-means clustering with that of deep modularization in terms of speech separation quality. Since $k$-means uses a frozen pre-trained encoder for feature extraction, we benchmarked its performance against MOD1 and MOD2, which also leverage a frozen encoder.

As shown in Table 7, deep modularization consistently outperformed $k$-means across all test datasets. These results suggest that deep modularization is more effective at capturing the relationships between T-F bins patches or raw speech blocks, resulting in superior mask generation and better speech separation quality compared to the classical $k$-means approach.