Multi-Scale Group Transformer for Long Sequence Modeling in Speech Separation

1 Introduction

Speech separation is a fundamental task in acoustic signal processing with a wide range of applications [Wang and Chen, 2018]. The goal of speech separation is to separate target speech from interfering speech, non-speech noise, or both. Thanks to the success of deep learning, the performance of single-channel speech separation system have been dramatically improved in recent years [Hershey et al., 2016; Kolbæk et al., 2017]. In particular, a new category of speech separation methods called time-domain methods [Luo and Mesgarani, 2018; Luo and Mesgarani, 2019] begin to emerge. These methods take the sampled data points from raw waveform as input, use a learnable neural network layer as encoder, and adopt a sequence modeling tool for feature separation before the separated feature is converted back to time domain by a learnable decoder. The time-domain methods have been shown to surpass ideal time-frequency (T-F) magnitude masking methods for speech separation.

∗This work is done when Yucheng Zhao is an intern in MSRA

†Corresponding author

Recently, several investigations [Heitkaemper et al., 2019] reveal that the performance of time-domain methods is highly dependent on using short frame length in the encoder, which produces extremely long sequence. For example, the frame length used in Conv-TasNet [Luo and Mesgarani, 2019] is 2ms, producing 4000 frames for a 4-second audio segment with 1ms overlap. Such a long sequence is difficult to model using the conventional recurrent neural network (RNN) or temporal convolutional network (TCN)[Bai et al., 2018] as both of them have a long path before connecting all positions [Vaswani et al., 2017].

Transformer [Vaswani et al., 2017] have recently shown its strong capability for sequence modeling in many natural language processing tasks [Devlin et al., 2018]. Compared to RNN or TCN, Transformer uses self-attention (SA) to compute correlations between any input positions in parallel, which can effectively capture global dependencies in addition to the local dependencies. Introducing Transformer to speech separation has the potential to model the long-range dependencies displayed in speech signals. But the main obstacle is the complexity. The computational complexity of Transformer grows quadratically with the sequence length, which is not affordable when the sequence length is on the order of thousands.

In order to solve this issue, we propose a novel architecture called multi-scale group Transformer (MSGT). We divide the input sequence into groups and use group self-attention to calculate correlations within each group. The complexity remains tractable as long as the group size is not too large. As we down-sample the features, the group size in group Transformer is kept constant, so correlations in a longer range can be captured. On the lowest scale, all data points are contained in a single group. As such, both local and global dependencies are captured and retained. Besides, this structure pays more attention to local dependencies, which aligns with the physical characteristics of acoustic signals.

We implement two versions of MSGT with different complexities and apply them to the Conv-TasNet speech separation system. In particular, we replace TCN with MSGT for sequence modeling. Experimental results show that the MSGT-based system, which we call MSGT-TasNet, achieves significant performance gain over Conv-TasNet in both noisefree speech separation on WSJ0-2mix dataset [Hershey et al., 2016] and noisy speech separation on WHAM! dataset [Wichern et al., 2019].

To summarize, our work makes three main contributions: (1) To the best of our knowledge, we are the first to investigate Transformer for sequence modeling in speech separation. (2) We propose multi-scale group Transformer which reduces the complexity of the standard Transformer without losing its capability to model global dependencies. (3) We implement two versions of MSGT and use them to build a speech separation system MSGT-TasNet. Without bells and whistles, the performance of MSGT-TasNet is already on par with the state-of-the-art (SOTA) methods.

2 RelatedWork

In this section, we provide the background of the speech separation task and briefly review several variations of Transformer for long sequence modeling.

2.1 Speech Separation

The goal of speech separation is to separate target speech from interference including interfering speech, non-speech noise, or both. As shown in Fig.2, a deep learning-based speech separation system is composed of three modules. Before the pioneering work named TasNet [Luo and Mesgarani, 2018], researchers have been using fixed transformation, such as STFT and ISTFT, for the encoder and the decoder [Hershey et al., 2016; Kolbæk et al., 2017; Shi et al., 2018]. The separation is conducted on the two-dimensional time-frequency features. However, the success of TasNet has ignited the interest of time-domain approaches, which directly take data samples from time-domain raw waveform as input, and conduct separation in a latent one-dimensional domain.

One typical feature of time-domain approaches is that their success rely on the effectiveness in modeling long sequences [Heitkaemper et al., 2019; Luo and Mesgarani,

2019]. For Conv-TasNet, it is discovered that shorter frame length achieves better performance than longer ones. Therefore, a frame length of 2ms is suggested for high performance. This length is an order of magnitude smaller than the length used by conventional encoders. Given a fixed-length audio input, such a short frame length also produces an extremely long sequence, which is difficult to model.

In Conv-TasNet [Luo and Mesgarani, 2019], the authors used a temporal convolution network to model long-range dependencies. However, it has been shown [Vaswani et al., 2017] that convolution is not as efficient as Transformer in sequence modeling, as the latter is capable of capturing both local and global dependencies. This has been the motivation of our work. More recently, FurcaNeXt [Zhang et al., 2020] introduced gated activation and ensemble learning into the framework to improve performance. However, they are still using TCN for sequence modeling.

2.2 Transformer

Transformer is a strong sequence modeling tool but it is not readily applicable to long sequences due to its quadratically growing computational complexity with the sequence length. Several pieces of work [Dai et al., 2019; Liu and Lapata, 2019; Shen et al., 2018; Miculicich et al., 2018] in natural language processing area have tried to tackle this limitation, so that Transformer can be applied to document-level tasks or having a higher efficiency. The main idea is to divide the sequence into a few conceptually meaningful sets or blocks, such as sentences and paragraphs, and then adopt a hierarchical architecture to explore the intra-set and inter-set correlations. To be more specific, local dependencies are captured within each set, and global dependencies are computed through cross-set attention.

[Shen et al., 2018] presents a bi-directional block selfattention network (SAN) that divides a sequence into blocks and sequentially applies intra-block SAN to each block and inter-block SAN across blocks. [Miculicich et al., 2018] uses a hierarchical attention network structure for document-level machine translation and [Liu and Lapata, 2019] uses a hierarchical Transformer for multi-document summarization. Transformer-XL [Dai et al., 2019] introduces a segment-level recurrence mechanism that enables Transformer to learn dependencies beyond a fixed length.

The main difference between these hierarchical SAN and our work is that we do not calculate set-level dependencies. NLP tasks involve semantic units such as sentence, paragraph, and document. However, speech separation is a lowlevel task without multiple semantic units. For such input data, our model uses multi-scale fusion architecture to calculate element-level dependencies at different resolutions.

There is also a work called Sparse Transformer [Child et al., 2019] for long sequence generation. They factorized full attention matrix by some sparse attention matrices to reduce the complexity. Compared to our model, Sparse Transformer relies on a highly optimized sparse matrix implementation and the complexity of it is O(n√n), which is higher than ours for long sequences.

3 Multi-Scale Group Transformer

We intend to design a long-sequence modeling tool based on Transformer. There are two design objectives. First, we shall reduce the computational complexity of Transformer so that it can scale with sequence length. Second, we shall maintain Transformer's ability to model both short-term and long-term dependencies.

3.1 The Proposed Architecture

Fig. 1 presents the architecture of the proposed multi-scale group Transformer (MSGT). The key innovations are the group self-attention and the multi-scale fusion.

In the original Transformer design, correlations are computed between any two positions in the input sequence. This brings quadratic complexity with respect to the sequence length. In contrast, the propose group self-attention restricts the correlation computation within local regions. When the group size is fixed at a constant, the number of groups grows linearly with the sequence length, so does the computational complexity of group self-attention.

However, group self-attention (GSA) does not consider correlations across groups, losing the capability to capture global dependencies. To retain this important capability and to balance the computation resources allocated between local and global dependencies, we propose the multi-scale architecture. On the high-resolution (large) scale, GSA captures local dependencies, while on the low-resolution (small) scale, GSA captures long-range dependencies. As sequences in small scales are several times shorter than sequences in large scales, more computation resources are thus allocated to capture local dependencies. This is consistent with the intuition that local correlations are stronger and more important than global correlations in audio signals.

We implement MSGT using two basic modules: operation module and transition module. The operation module conducts multi-scale feature transformation. It may conduct different transformations, including group Transformer and identity mapping, on different scales. The transition module handles feature resizing and feature fusion.

Group Self-Attention

Given an input sequence X ∈ Rn×d where d is the feature dimension and n is the sequence length, GSA first divides the sequence into s equal-sized non-overlapped groups with group size g. Then, standard self-attention is calculated within each group and the output of each group is concatenated. We formulate the above process as follows:

The standard self-attention, which proposed in [Vaswani et al., 2017] is computed as follows:

where √d is the scaling factor, W Q,W K ,W V ∈ Rd×d are parameter matrices of different linear transformations.

For simplicity, here we only describe the formula of singlehead group self-attention. The extension to multi-head group self-attention is straightforward.

Multi-Scale Fusion

Multi-scale fusion is conducted by the transition module. It exchanges information across multi-scale features. Let us take 3 scales as a example, which is illustrated in Figure 4. For the largest scale, we first conduct up-sampling on two smaller scales. Note we gradually up-sample the smallest scale by using two sequential up-sampling networks. Then, we have three features with the same size and use addition to fuse them. For the medium scale, we up-sample the smallest scale, down-sample the largest scale and then add them together. The operation on the smallest scale is similar. Note that we do not have to fuse all scales in the transition module. It is a design choice how many scales are fused. We implement two versions of MSGT with different complexities.

3.2 Two Implementations of MSGT

We present two concrete design instances of MSGT, namely dense-fusion MSGT and light-fusion MSGT, in Fig. 3. The fusion structures of these two implementations are inspired by the HRNet [Sun et al., 2019] and the U-Net [Ronneberger et al., 2015] in the computer vision area.

Dense-Fusion MSGT

The architecture of dense-fusion MSGT is illustrated in Fig. 3a. With dense fusion, multiple group Transformers are applied on each scale. For simplicity, we fuse only the adjacent scale features in the transition module. In other words, an input to the group Transformer is fused from no more than three scales. The dense fusion has log-linear complexity with sequence length.

Light-Fusion MSGT

As illustrated in Fig. 3b, the light fusion version only use one GT on each scale, although each GT may contain multiple layers of transformation. The transition module does not do fusion as the features scale down. When the features scale up, only features from two adjacent scales are fused together. The light fusion has linear complexity with sequence length.

3.3 Analysis

The sequence modeling problem can be formulated as mapping one sequence {x1, x2, ..., xn} to another sequence of equal length {z1, z2, ..., zn}, where xi, zi ∈ Rd denote the feature vector of a symbol in the sequence.

There are three desired properties for a sequence modeling tool: 1) Parallel computing; 2) Capability to capture longrange dependencies; 3) Low-order complexity which allows it to scale to long sequences. The original Transformer paper [Vaswani et al., 2017] introduced three metrics to measure these desired properties. They are complexity per layer, minimum number of sequential operation required and the maximum path length between any two input and output position.

Table 1 provides a comparison of the tree metrics between the proposed MSGT and other widely used sequence modeling tools. A standard self-attention layer has complexity of O(n2 · d), which is too high for long sequences. The multi-scale group self-attention have maximum path length O(log2(n/g)) and complexity O(g · n · d) or O(g · n · d · log (n/g)) depending on the fusion choices. We will empirically show this improvement in Section 5.4. For completeness, we also include recurrent and convolutional layers in this table, but they are inferior to self-attention which has been detailed in [Vaswani et al., 2017].

4 MSGT for Speech Separation

The goal of speech separation is to separate target speech from a mixture signal. Following the convention, all the speech signals appeared in the mixture are treated as target speeches. The mixture can be denoted by a sequence y(t) in the time domain. It can be decomposed as the summation of K speech signals xk(t), and an additive noise n(t),

where t is the time index.

Recent deep learning-based speech separation systems consist of three main components, namely the encoder, the decoder, and the separator. Fig.2 illustrates this framework. The encoder processes time domain mixture signal y(t) by first dividing it into n overlapping frames yj ∈ R1×L (j = 1...n), where L denotes frame length, and j is the frame index. Then, each frame is transformed into d-dimensional features Yj ∈ R1×d . Features of all the n frames are concatenated to form Y ∈ Rn×d . The separator use some sequence modeling tool to estimate multiplicative masks Mk ∈ Rn×d for each signal and then multiply it on encoded features Y , producing ˆXk as separated features. Finally, the decoder transform ˆXk back to time domain and output the separated signals ˆxk(t).

It is worth noting that it is not necessary for the separated signals to have the same permutation as the ground-truth label. Utterance-level permutation invariant training (uPIT) [Kolbæk et al., 2017] can be applied to address this problem.

We adopt the well-known speech separation framework Conv-TasNet [Luo and Mesgarani, 2019] and replace the original TCN with the proposed MSGT in its separator. The encoder and the decoder are kept the same. Our system is still a time-domain method, so we name it MSGTTasNet. We train MSTG-TasNet using the time domain scaleinvariant signal-to-distortion ratio (SI-SDR) as the training target, which is defined as:

where ˆ s ∈ RT and s ∈ RT are the estimated and original clean signals, respectively. α = ˆ sT s/|s|2 is the scaling factor.

5 Experiments

5.1 Dataset and EvaluationMetric

We use two datasets for evaluation. The first is the widelyused WSJ0-2mix dataset [Hershey et al., 2016]. WSJ0-2mix contains 30 hours of training data, 10 hours of validation data and 5 hours of testing data. The validation set contains utterances of the same speakers as in the training set, while the testing set contains utterances of different speakers. Each mixture is artificially generated at a random signal-to-noise ratio (SNR) between -5 and 5 dB. The second is the recently proposed WHAM! dataset [Wichern et al., 2019], which is an extension of the original WSJ0-2mix. The WHAM! dataset consists of two speaker mixtures from the WSJ0-2mix dataset combined with real ambient noise samples. This is a more challenging dataset compared to the noise-free WSJ0-2mix.

We use the scale-invariant signal-to-distortion ratio (SISDR) improvement [Le Roux et al., 2019] and signal-todistortion ratio (SDR) [Vincent et al., 2006] improvement as the main evaluation metrics. SI-SDR is also referred to as SISNR in some work [Luo and Mesgarani, 2019]. We also report perceptual evaluation of subjective quality (PESQ) [Rix et al., 2001] to evaluate the quality of the separated mixtures.

5.2 Experiment Configuration

We train all models for 1M steps on 4-second segments with sample rate of 8K Hz. We use Adam optimizer with warmup. The learning rate is initialized to 0.0003 and is adjusted according to the following formula:

We choose warmup steps = 10000. We also use dropout to relieve over-fitting.

In all experiments, we use frame length of 2 ms. We choose group size of 1000 for noise-free speech separation and group size of 500 for noisy speech separation. Following the notations in [Vaswani et al., 2017], the Transformer parameters are dff = 1024, dmodel = 512, and h = 8. In the light fusion, we use 8 layers of transformation for the GT in the smallest scale and 2 layers for the GT in the other scales. In the dense fusion, we use 3 layers of transformation for GT in all the scales. The output feature dimension of the encoder is 1024.

5.3 Comparison with Other Methods

Previously, research on speech separation is more focused on separating target speech from interfering speech [Luo and Mesgarani, 2019; Liu and Wang, 2019]. In this paper, we use the term noise-free speech separation to denote this conventional setting while we use another term noisy speech separation to refer to the task of separating target speech from both interfering speech and non-speech noise. The WSJ0-2mix dataset is used for the former task while the WHAM! dataset is used for the latter task.

Speech Separation Results on WSJ0-2mix

Table 2 shows speech separation results on WSJ0-2mix. DPCL++ [Isik et al., 2016], uPIT [Kolbæk et al., 2017], chimera++ [Wang et al., 2018] and Sing Prediction Net [Wang et al., 2019] are time-frequency (T-F) domain methods. Conv-TasNet [Luo and Mesgarani, 2019], FurcaNeXt [Zhang et al., 2020] and our proposed MSG-TasNet are timedomain methods. The fist row gives the performance of ideal phase-sensitive mask (PSM), which is a reasonable upper bound for all T-F domain methods.

Our proposed two versions of MSGT-TasNet outperform Conv-TasNet by a large margin only by replacing TCN with MSGT, showing that MSGT is a better sequence modeling method for speech separation. The SI-SDRi performance of the light version is 0.2dB lower than that of the dense version, but it uses 2x fewer memory and runs 2.8x faster.

FurcaNeXt and Deep CASA achieve slightly higher SDRi or SI-SDRi than MSGT-TasNet. But the innovations in these two methods can also be applied to our framework and can potentially improve the performance of MSGT-TasNet. Specifically, FurcaNext adopts gated activation and Deep CASA optimizes frame-level separation and speaker tracking in turn. We plan to integrate these two innovations into our model in the future.

Noisy Speech Separation Results on WHAM!

The WHAM! dataset is more realistic and challenging since it contains noise. As this is a new dataset and only baseline result is available, we add several idea masks to the comparison. The ideal masks provide performance upper bounds for most T-F domain methods. Results are presented in Table 3.

In noisy speech separation, MSGT-TasNet (Light) and MSGT-TasNet (Dense) get 0.3 dB and 1.1 dB SI-SDRi gain over Conv-TasNet, respectively. This gain is smaller than what we get in noise-free speech separation, but the advantage of the dense version to the light version becomes significant. Both results are due to the fact that noisy speech separation is a more difficult task than its noise-free counterpart. Note that, in this difficult task, MSGT-TasNet (Dense) is the first method that surpasses the ideal T-F mask IRM.

a) GPU Memory Cost vs Sequence Length

Figure 5: Empirical computation cost of Transformer and MSGT. D and L stand for dense and light, respectively. The number follows indicates group size. One sample is a 4s audio segment.

5.4 Efficiency for Long Sequence

MSGT reduces the time/space complexity from O(n2 · d) to O(g · n · d) or O(g · n · d · log (n/g)) in theory. We would like to compare the empirical GPU memory cost and throughput for speech separation under the same configuration. We use 4s audio segment and choose different frame lengths to control the actual sequence length used in group Transformer. We choose group size g as a hyper-parameter. For different sequence lengths, we keep the total number of group Transformers the same for light fusion and we keep the number of group Transformer in the largest scale the same for dense fusion. This constraint yields models with similar size for different sequence length.

As Fig.5a shows, the GPU memory usage in MSGT is significantly smaller than in Transformer, and it grows slower with the increase of sequence length. Under the same group size, the light version requires less GPU memory than the dense version. Within the same fusion model, the GPU memory cost grows with the group sizes g.

Fig.5b presents the throughput of different models. The numbers are acquired on a single P100 GPU. MSGT (Dense) has slightly lower throughput than Transformer for short sequence because more scales are involved in computation. MSGT (Light) consistently achieves higher throughput than Transformer.

5.5 Ablation Study

All ablation studies are carried out on the WSJ0-2mix speech separation dataset.

Comparison of different sequence modeling tools. As shown in Table 4, using group Transformer without multiscale fusion significantly degrades the performance. The performance of single-scale GT, which has only one scale and uses the same number of layers as MSGT (Light), is even worse than TCN. Using light fusion, we can achieve almost doubled throughput and 3.3dB SI-SDRi gain over the singlescale version. We would also like to mention that MSGT (Light) achieve 1.5 dB gain over TCN with 3.9X speedup.

Influence of group size. The hyper-parameter g in group Transformer controls how many frames should be consider as a group. As shown in Fig.6, the largest group size achieves the highest SI-SDRi for both fusion versions. In particular, large group size is essential for the light version to get good performance, which may due to its inadequacy in exploring long-range dependencies.

6 Conclusion and FutureWork

In this paper, we present the design and implementation of multi-scale group Transformer for long sequence modeling. Through group self-attention and multi-scale fusion, MSGT significantly reduces the computational complexity of Transformer without affecting its performance. Two versions of MSGT with different complexities are implemented and applied in a well-known speech separation framework called Conv-TasNet. Experiment results show the proposed MSGT-TasNet achieves a large gain over Conv-TasNet on both WSJ0-2min and WHAM! benchmarks. For the noisy speech separation task, MSGT-TasNet is the first approach surpassing the performance of the ideal T-F mask.

In the future, we plan to apply neural architecture search (NAS) [Liu et al., 2018] to find better fusion choices of MSGT. Besides, we plan to integrate the innovations of Deep CASA into MSGT-TasNet. We believe that such combination can further advance the state-of-the-art of speech separation.