{'fraction_non_alphanumeric': 0.042638134781220055, 'fraction_numerical': 0.016022569998935376, 'mean_word_length': 4.514235397710596, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Diffusion-based Generative Speech Source Separation

We propose DiffSep, a new single channel source separation method based on score-matching of a stochastic differential equation (SDE). We craft a tailored continuous time diffusion-mixing process starting from the separated sources and converging to a Gaussian distribution centered on their mixture. This formulation lets us apply the machinery of score-based generative modelling. First, we train a neural network to approximate the score function of the marginal probabilities or the diffusion-mixing process. Then, we use it to solve the reverse time SDE that progressively separates the sources starting from their mixture. We propose a modified training strategy to handle model mismatch and source permutation ambiguity. Experiments on the WSJ0 2mix dataset demonstrate the potential of the method. Furthermore, the method is also suitable for speech enhancement and shows performance competitive with prior work on the VoiceBank-DEMAND dataset.


INTRODUCTION
Source separation (SS) refers to a family of techniques that can be used to recover signals of interests from their mixtures. It has broad applications, but we focus on single channel audio and speech [1]. This is an underdetermined problem where several signals are recovered from a single mixture. Early success in this area was made with methods such as non-negative matrix factorization (NMF) [2]. Impressive progress was brought by data-driven methods and deep neural networks (DNN). We mention deep clustering where a DNN is trained to create an embedding space in which the time-frequency bins of the input spectrogram cluster by source [3]. Then, one can recover isolated sources by creating masks corresponding to clusters of bins and multiply the spectrogram with it. These methods then evolved to networks predicting mask values directly [4]. All these data-driven methods require to solve the fundamental source permutation ambiguity. Namely, there is no inherent preferred order for the sources. This is usually handled by permutation invariant training (PIT) [5]. In PIT, the objective function is computed for all permutations of the sources and the one with minimum value is chosen to compute the gradient. Recently very strong time-domain baselines such as Conv-TasNet [6] or dual-path transformer [7] have been proposed. We mention also all-attention based [8], multiscale networks [9], WaveSplit [10], and the recent state-of-the-art TF-GridNet [11] back in the time-frequency domain. For a survey and benchmarks of recent methods see [11], [12]. All these recent approaches are trained with a discriminative objective such as the scale-invariant signal-to-distortion ratio (SI-SDR) [13]. Except the Code will be shared after the review process. traditional methods, such as NMF, few approaches leverage generative models. One exception adversarially trains speaker-specific generative models and uses them for separation [14].
Unlike discriminative methods, generative modelling allows to approximate complex data distributions [15]. Recently, score-based generative modelling (SGM) [16], also known as diffusion-based modelling, has made rapid and impressive progress, in particular in the domain of image generation [17]. SGM defines a forward process that progressively transforms a sample of the target distribution into Gaussian noise. Given its score-function, i.e., the gradient of the log-probability, the target distribution can be sampled by running a reverse process starting from noise. Crucially, the score-function is (usually) unknown, but can be approximated by a DNN using a simple training strategy [18]. Recent approaches result from the combination of score-matching and Langevin sampling [16], [19], [20]. Two frameworks have been proposed to unify various approaches, one based on graphical models [21], and the other on stochastic differential equations (SDE) [22]. SGM has quickly found application in audio and speech processing. It has been successfully applied in speech synthesis [23]- [26], speech restoration [27], [28], and bandwidth extension [29]. Surprisingly, perhaps, several SGMs have been proposed for speech enhancement, a typically discriminative task [30]- [33]. Besides sound, but perhaps closest to this work, separation of images using deep generative priors and Langevin sampling has been demonstrated [34]. However, to the best of our knowledge, a diffusion-based speech separation method has yet to be proposed.
We fill this gap by proposing what we believe is the first diffusion-based approach for the separation of speech signals. The proposed algorithm, DiffSep, is based on a carefully designed SDE that starts from the separated signals and converges in average to their mixture. Then, by solving the corresponding reverse-time SDE, it is possible to recover the individual sources from their mixture. See Fig. 1 for an illustration. We propose a modified training strategy for the score-matching network that addresses ambiguity in the assignment of the sources at the beginning of the inference process. We demonstrate the method on the widely used WSJ0 2mix benchmark dataset [3]. While it does not outperform, yet, discriminative baselines in terms of separation metrics, it achieves high non-intrusive DNSMOS P.835 [35] score, indicating naturalness. In addition, our approach is also suitable for speech enhancement by considering the noise as just an extra source, and performs comparably to prior work [32] on the VoiceBank-DEMAND dataset [36]. We represent audio signals with N samples by real valued vectors from R N . We introduce a simple mixture model with K sources, where s k ∈ R N for all k. We remark that this model can account for environmental noise, reverberation, and other degradation by adding an extra source. By concatenating all the source signals, we obtain vectors in R KN . We define in this notation the vector of separated sources and their average value, We can define the time-invariant mixing operation as a multiplication by (A ⊗ IN ) where A is a K × K matrix, and ⊗ is the Kronecker product. Multiplying a vector v ∈ R KN by such a matrix amounts to multiplying by A all sub-vectors of length K with elements taken from v at interval N , i.e., To lighten the notation, we overload the regular matrix product as Av (A ⊗ IN )v in the rest of this paper. Next, we define the projection matrices where 1 is the all one vector, of size K here. The matrix P projects onto the subspace of average values, andP its orthogonal complement. For example, a compact alternative to (1) is P s =s.

SDE for Score-based Generative Modelling
Score-based generative modelling (SGM) via SDE is a principled approach to model complex data distributions [22]. First, a diffusion process starting from samples from the target distribution and converging to a Gaussian distribution can be described by the following stochastic differential equation where xt : R N → R N is a vector-valued function of the time t ∈ R, dxt is its derivative with respect to t, and dt an infinitesimal time step. Functions f : R N → R N and g : R → R are called the drift and diffusion coefficients of xt. The term dwt is a standard Wiener process. For more details on SDEs, see [37]. We insist that the time t in the diffusion process is unrelated to the time of the audio signals that will appear later in this paper. A remarkable property of (5) is that under some mild conditions, there exists a corresponding reverse-time SDE, going from t = T to t = 0 [38]. Here, dw is a reverse Brownian process, and pt(x) is the marginal distribution of xt.
During the forward process, x0 being known, it is usually possible to have a closed-form formula for pt(x). For an affine f (x, t), it is Gaussian and its parameters, i.e., mean and covariance matrix, usually tractable [37]. However, in the reverse process, pt(x) is unknown, and thus (6) cannot be solved directly. The key idea of SGM [16], [19], [22] is to train a neural network q θ (x, t) so that q θ (x, t) ≈ ∇x log pt(x). Provided that the approximation is good enough, we can generate samples from the distribution of x0 by numerically solving (6), replacing ∇x log pt(x) by q θ (x, t). Fig. 1: Illustration of the diffusion-mixing process. From left to right, the sources start separated and progressively mix as more and more noise is added. The reverse process starts from the mixture and ends with separated sources. The spectrograms were produced by actual separation of the mixture on the right using the proposed method.

Related Work in Diffusion-based Speech Processing
SGM first found widespread adoption for the generation of very high quality synthetic images, e.g. [17]. Speech synthesis is thus a logical candidate for its adoption [23], [24]. In this area, several approaches that shape the noise adaptively for the target speech have been proposed. One uses the local signal energy [25], while the other introduces correlation in the noise according to the target spectrogram [26]. Bandwidth extension and restoration also require to generate missing parts of the audio and high quality diffusion-based approaches have been proposed [27]- [29]. Diffuse [30] and CDiffuse [31] use conditional probabilistic diffusion to enhance a noisy speech signal. Also for speech enhancement, Richter et al. have proposed another approach [32], [33] based on solving the reverse-time SDE corresponding to the forward model This is the Ornstein-Uhlenbeck SDE that, in the forward direction, takes the clean speech s progressively towards the noisy one y.

DIFFUSION-BASED SOURCE SEPARATION
We propose a framework for applying SGM to the source separation problem. Essentially, we design a forward SDE where both diffusion and mixing across sources happen over time. Each step of the process can be described as adding an infinitesimal amount of noise and doing an infinitesimal amount of mixing. This can be formalized as the following SDE whereP is defined in (4) and s in (2). We use the diffusion coefficient of the Variance Exploding SDE of Song et al. [22], The SDE (8) has the interesting property that the mean of the marginal µ t goes from the vector of separated signals at t = 0 to the mixture vectors as t grows large. This is formalized in the following theorem proved in Appendix 6.
where if we let ξ1 = 0, ξ2 = γ, and ρ = σmax σ min , then, Fig. 1 shows an example of the process on a real mixture. We now have an explicit expression for a sample xt, where Lt = λ1(t) 1/2 P + λ2(t) 1/2P . Fig. 2 shows the evolution of some of the parameters of the distribution of xt over time. We can make the difference of µ T ands arbitrarily small by tuning the value of γ. In addition, we observe that the correlation of the noise added onto both sources increases over time due to the mixing process.
Previous work on speech enhancement has successfully applied the diffusion process in the short-time Fourier transform (STFT) domain with a non-linear transform [32], [33]. The approach we propose could be equally applied in the time-domain or in any linearly transformed domain such as the STFT. However, since the process of (8) models the linear mixing of the sources, we cannot apply a non-linear transformation. Instead, we perform the diffusion process in the time-domain, but operate our network in the non-linear STFT domain of [33], as described in Section 4.2.

Inference
At inference time, the separation is done by solving (6) with the help of a score-matching network q θ (x, t, y) ≈ ∇x log pt(x). The initial value of the process is sampled from the distribution, Then, we adopt the predictor-corrector approach described in [22] to solve (6). The prediction step is done by reverse diffuse sampling. The correction is done by annealed Langevin sampling [20].

Permutation and Mismatch-aware Training Procedure
To train the score network q θ (x, t, y), we propose a slightly modified score-matching procedure. But, first, we follow the usual procedure. Thanks to Theorem 1, the score function has a closed form expression. The computation is similar to that in [32], and we obtain The training loss is where z ∼ N (0, IKN ), t ∼ U(t , T ), and x0 is chosen at random from the dataset. We use the norm induced by Σt which is equivalent to the weighting scheme proposed by Song et al. [16]. This scheme makes the order of magnitude of the cost independent of Σt and lets us avoid the computation of the inverse of Lt. However, when training only with the procedure just described, we found poor performance at inference time. We identified two problems. First, there is a discrepancy between E[xT ] =s from (14) and µ T . While µ T contains slightly different ratios of the sources in its components,s does not. Second, the network needs to decide in which order to output the sources. This is usually solved with a PIT objective [5]. To solve these challenges, we propose to train the network in a way that includes the model mismatch.
For each sample, with probability 1 − pT , pT ∈ [0, 1], we apply the regular score-matching procedure minimizing (17). With probability pT , we set t = T and minimize the following alternative loss, wherexT is sampled from (14), P is the set of permutations of the sources, and µ T (π) is µ T computed for the permutation π. The rational for this objective is that the score network will learn to remove both the noise and the model mismatch. We note that for speech enhancement, there is no ambiguity in the order of sources and the minimization over permutations can be removed.

Datasets
We use the WSJ0 2mix dataset [3] for the speech separation experiment. It contains 20000 (30 h), 5000 (10 h), and 3000 (5 h) two speakers mixtures whose individual speech samples were taken from the WSJ0 dataset. The utterances are fully overlapped and speakers do not overlap between training and test datasets. The relative signal-to-noise ratio (SNR) of the mixtures is from −5 dB to 5 dB. The sampling frequency is 8 kHz. In addition, we use the clean test set of Libri2Mix [39] to test mismatched conditions. For the speech enhancement task, we use the VoiceBank-DEMAND dataset [36]. It is a collection of clean speech samples of 30 speakers from the VoiceBank corpus degraded by noise, recorded as well as artificial. The dataset is split into train (8.6 h), validation (0.7 h), and test (0.6 h). The SNR of the mixtures ranges from 0 dB to 15 dB for training and validation, and from 2.5 dB to 17.5 dB for testing. The sampling rate of the mixtures is 16 kHz.

Model
We use the noise conditioned score-matching network (NCSN++) of Song et al. [22]. An STFT and its inverse layer sandwich the network and, in addition, we apply the non-linear c(x) transform proposed in [33] after the STFT, and its inverse before the iSTFT, The real and imaginary parts are concatenated. The network has K + 1 inputs, one for each of the target sources plus one for the mixture on which we condition, and K outputs. In the low quality sample, the spectrogram looks good, but we can see block permutations of the sources occurring.

Hyperparameters for Inference and Training
For inference, we use 30 prediction steps. For each one, we do one correction step with step size r = 0.5. The networks are trained using Pytorch with the Adam optimizer. We use exponential averaging of the network weights as in [32]. Separation We use γ = 2, σmin = 0.05, σmax = 0.5. For the transformation (18), we use α = 0.5 and β = 0.15. The probability of selecting a sample for modified training is pT = 0.1. The effective batch size is 48 and the learning rate 0.0005. We trained for 1000 epochs and we chose the one with largest SI-SDR on a sentinel validation batch.
Enhancement The SDE parameters are γ = 2, σmin = 0.05, σmax = 0.5. The transformation (18) uses α = 0.5 and β = 0.15. We found it beneficial to use PriorGrad [25] for the noise generation. The variance of the process noise for a given sample is chosen as the average power of the 500 closest samples in the mixture. In the enhancement setting, there is no permutation ambiguity so we always order the sources with speech first and noise second. We use pT = 0.03. The effective batch size was 16 and the learning rate 0.0001. The network was trained for around 160 epochs.
Separation Table 1 shows the results of the speech separation experiments. We also train a baseline Conv-TasNet [6] network. On the intrusive metrics, SI-SDR, PESQ, and STOI, the proposed method scores lower than the baseline. However, it scores higher on the non-intrusive OVRL metric. Informal listening tests further reveal that low separation metrics are often due to very natural sounding permutations of the speakers. This problem seems to happen for similar sounding speakers. Example spectrograms of poor, fair, and high quality separation results are shown in Fig. 3. It was reported that diffusion-based enhancement performs well in mismatched condition [32]. We thus evaluated the model trained on WSJ0 2mix on the Libri2Mix [39] clean test set. The proposed method still scores lower than Conv-TasNet for SI-SDR, PESQ, and STOI, but the metric degradation is smaller, e.g. 1 dB less reduction of SI-SDR. On the other hand, the proposed method still enjoys fairly high OVRL score, indicating good perceptual quality.
Enhancement Table 2 shows the results of the evaluation of the noisy speech enhancement task. The proposed method performs very similarly to other diffusion methods CDiffuse [31] and SGMSE+ [32]. We obtain an SI-SDR higher by 0.2 dB, but PESQ value closer to that of CDiffuse. ESTOI is somewhere in the middle. The OVRL metric suggests good perceptual quality, but was not provided for the other methods. In terms of SI-SDR particularly, a gap remains with discriminative methods such as Conv-TasNet.

CONCLUSION
We proposed DiffSep, the first speech source separation method based on diffusion process. We successfully separated speech sources, and demonstrated that our approach extends to noisy speech enhancement. For separation, there was a substantial gap in performance with state-of-the-art discriminative methods. We believe this may be due to both outputs being from the same distribution, thus confusing the generative model. We have done initial work to alleviate the permutation ambiguity, using a modified training strategy for the score-matching network, but further work is likely necessary to robustify the inference procedure against permutation mistakes. Bringing in proven architectures from the separation literature is an interesting area of future work. We have also not used yet any of the theoretical properties of the distribution of mixtures, which could be elegantly integrated in the score network architecture, as in [34].