Tiny-Sepformer: A Tiny Time-Domain Transformer Network for Speech Separation

The overall model architecture is depicted in Figure 1. The time-domain mixture-signal $X$ is firstly fed into the encoder $\theta_{enc}$. The encoder $\theta_{enc}$ is a single 1D-convolutional layer, followed by ReLU activation. The hidden feature representation $H$ is extracted through the encoder.

$$H=\theta_{enc}(X)=\mathrm{ReLU}(\mathrm{Conv1D}(X))$$

(1)

After that, the encoder output $H$ is used as the input of the masking network $\theta_{mask}$, to produce mask matrices $M_{1:K}$.

$$M_{1:K}=\theta_{mask}(H)$$

(2)

Finally, the mask matrices $M_{1:K}$ and hidden features of encoder $H$ are as the input of the decoder $\theta_{dec}$. The decoder $\theta_{dec}$ is a transposed convolutional layer with the same stride and kernel size of the encoder $\theta_{enc}$. The decoder outputs the separated signal $\hat{X}_{1:K}$ from each source $\left\{1,2,...,K\right\}$.

$$\hat{X}_{k}=\theta_{dec}(M_{k},H)=\mathrm{Conv1D^{T}}(M_{k}\odot H)$$

(3)

In which, $\mathrm{Conv1D^{T}}$ is the transposed convolution, and $\odot$ is denoted as element-wise multiplication. The objective of model training is maximizing the Scale-Invariant Source-to-Noise Ratio (SI-SNR). We use Utterance-level Permutation Invariant Training (uPIT) [24, 25] loss during training to deal with the label permutation problem.

The masking network consists of three steps: (1) pre-processing and chunking, (2) the proposed tiny Convolution-Attention (CA) blocks $\theta_{ca}$, and (3) post-processing and overlap-add.

Firstly, the encoder feature sequence $H$ is normalized with layer normalization, and processed by a linear layer to produce the $H^{d}$ with dimension $T\times D$.

$$H^{d}=\mathrm{Linear}(\mathrm{LayerNorm}(H))$$

(4)

After pre-processing, the hidden feature sequence $H^{d}$ is processed by the segmentation operation. The segmentation splits two dimensional $H^{d}\in T\times D$ into three dimensional chunks $H^{s}\in T_{S}\times S\times D$ with $50\%$ overlap. In which, $S$ is the chunk size of segmentation, and $T_{S}$ is the chunk number of segmentation result. The chunked $H^{s}$ is fed into several CA blocks $\theta_{ca}$, which will be detailed in Section 2.3. The CA blocks will be performed by $N_{mask}$ times iteratively to produce $H^{ca}$.

$$H^{ca}=\theta_{ca}(H^{s})_{\times N_{mask}}$$

(5)

The output of CA block $H^{ca}$ remains the same dimension $T_{S}\times S\times D$ with the input $H^{s}$. Then, $H^{ca}$ is processed by a linear layer with dimension $(D\times K)$ and with using the PReLU activation.

$$H^{dk}=\mathrm{PReLU}(\mathrm{Linear}(H^{ca}))$$

(6)

This post-processing operation generates feature maps $H^{dk}\in T_{S}\times S\times(D\times K)$ for each of the $K$ speakers. Afterwards, the Overlap-Add [13] is operated on $H^{dk}$, merging the three dimensional chunked sequence into two dimensinal feature sequence $H^{o}\in T\times(D\times K)$.

$$M_{1:K}=\mathrm{ReLU}(\mathrm{Linear}(H^{o})_{\times 2})$$

(7)

At the end, this representation $H^{o}$ is through two linear layers and a ReLU activation to obtain the mask matrices $M_{1:K}$ for each of the $K$ speakers. The dimension of each mask $M_{K}$ is $T\times D$.

The Convolution-Attention (CA) Block $\theta_{ca}$ is the core module of masking network. There are two kinds of network in the masking network: (1) IntraCA network $\theta_{ca}^{intra}$, operating within each chunk to model local features, (2) InterCA network $\theta_{ca}^{inter}$, processing between all the chunks to capture long dependency. The $\theta_{ca}$ block are designed to process $\theta_{ca}^{intra}$ network firstly, then permute the feature dimension from $T_{S}\times S\times D$ to $S\times T_{S}\times D$, and process $\theta_{ca}^{inter}$ network at last.

	$\displaystyle H^{ca}$	$\displaystyle=\theta_{ca}(H^{s})$		(8)
		$\displaystyle=\theta_{ca}^{inter}(\mathcal{P}(\theta_{ca}^{intra}(H^{s})_{\times N_{intra}}))_{\times N_{inter}}$

In which, $\mathcal{P}$ is denoted as the permutation operation. IntraCA $\theta_{ca}^{intra}$ is performed $N_{intra}$ times, and InterCA $\theta_{ca}^{inter}$ will perform by $N_{inter}$ times iteratively.

As depicted in Figure 2, IntraCA or InterCA network contains two parts: (1) multi-head attention, (2) 1D depthwise separable convolution. IntraCA $\theta_{ca}^{intra}$ and InterCA $\theta_{ca}^{inter}$ have the same network structure, but have different network parameters. Taking IntraCA $\theta_{ca}^{intra}$ as the example, the input feature $H^{s}\in T_{S}\times S\times D$ is split into two paths: (1) attention feature $H^{s}_{a}\in T_{S}\times S\times D_{a}$, and (2) convolution feature $H^{s}_{c}\in T_{S}\times S\times D_{c}$, where $D_{c}+D_{a}=D$. The choices of $D_{c}$ and $D_{a}$ will be analyzed in Section 3.4. The attention feature $H^{s}_{a}$ is fed into the standard multi-head attention mechanism, followed by layer normalization.

$$H^{a}=\mathrm{LayerNorm}(\mathrm{MultiHeadAttention}(H^{s}_{a})+H^{s}_{a})$$

(9)

Mean while, the other branch $H^{s}_{c}$ is processed by the 1D depthwise separable convolution as follows:

$$H^{s}_{d}=\mathrm{DepthwiseConv1D}(H^{s}_{c})$$

(10)

$$H^{c}_{p}=\mathrm{PointwiseConv1D}(H^{s}_{d})$$

(11)

followed by residual addition and layer normalization.

$$H^{c}=\mathrm{LayerNorm}(H^{c}_{p}+H^{s}_{c})$$

(12)

The 1D depthwise separable convolution could dramaticly reduce the number of network parameters. In addition, the dimension reduction of multi-head attention $D_{a}$ could also make a light attention layer.

After that, the output of two path $H^{a}$ and $H^{c}$ are concatenated together, feeding into feed-forward layers as follows:

$$H^{ca}_{f}=\mathrm{FeedForward}(H^{a}\oplus H^{c})$$

(13)

In which, $\oplus$ is denoted as concatenated operation. The dimension of feed-forward layer is denoted as $D_{f}$. Following the feed-forward layer are residual addition and layer normalization, similar to the vanilla Transformer network.

$$H^{ca}=\mathrm{LayerNorm}(H^{ca}_{f}+H^{a}\oplus H^{c})$$

(14)

Another idea of parameters reduction is cross-layer parameter sharing. There are several methods of sharing parameters, such as sharing the feed-forward network parameters, or only sharing attention and convolution parameters. In this work, we propose to share all the layer parameters within IntraCA/InterCA network, but different IntraCA/InterCA networks have different parameters. Therefore, with parameter sharing, the Equation 8 is modified as following:

$$H^{ca}=\theta_{ca}(H^{s})=\coprod^{N_{inter}}\theta_{ca}^{inter}(\mathcal{P}(\coprod^{N_{intra}}\theta_{ca}^{intra}(H^{s})))$$

(15)

In which, $\coprod^{N}$ means that performing the network layer by $N$ times iteratively, but each iteration has the same layer parameters. With this approach, the parameters of $\theta_{ca}^{intra}$ could decrease to ${1/N_{intra}}$ of the original network, as well as the ${1/N_{inter}}$ parameters for $\theta_{ca}^{inter}$. We share all the layers in IntraCA $\theta_{ca}^{intra}$ and InterCA $\theta_{ca}^{inter}$ respectively, resulting in much fewer network parameters of the Tiny-Sepformer model.

The most relevant structure to CA block is Conformer [19], which connects the attention and convolution in serial. The motivation of our work is to reduce the model parameters. Therefore, we design to connect the multi-head attention and 1D depthwise separable convolution in parallel. For simplicity, in this section, we denoted the input feature of the block as $H^{s}\in T\times D$, where $T$ is the number of time steps and $D$ is the channel dimension. Then, the parameters of attention and convolution paths are as Table 1:

The query, key, value and multi-head projection matrix has $D^{2}$ parameters respectively. Thus, for multi-head attention, the parameters are $4\times D^{2}$. The separable convolution has a depthwise convolutional layer with kernel size $K$ on each channel individually, and a pointwise convolutional layer on each frame independently but across all channels $D$. Therefore, for separable convolution, the parameters are $K\times D+D^{2}$. If connecting Attention+Convolution in Serial (Conformer) and keeping the feature dimension as $D$, their parameters are added together ($K\times D+5\times D^{2}$). For our CA block (Attention+Convolution in Parallel), the input feature is split into two branches: $D_{a}+D_{c}=D$. In Table 1, we set $D_{a}=D_{c}=\frac{1}{2}\times D$ for simplicity, then the parameters are reduced to $\frac{K}{2}\times D+\frac{5}{4}\times D^{2}$. Through the channel division operation, the parameters of CA block are much less than vanilla Transformer and Conformer, and still keep the feature dimension as $D$ into next feedforward layer.

We use the WSJ0-2mix and WSJ0-3mix [26] datasets to evaluate our Tiny-Sepformer model. These two datasets were generated by randomly selecting utterances from the WSJ0 corpus, and mixing them with two and three speakers. $30$ hours of training, $10$ hours of validation and $5$ hours of test speech dataset were used for all the experiments. All the speech were downsampled to $8$kHz in the data pre-processing.

We conducted all the experiments using the SpeechBrain toolkit [27]. For encoder $\theta_{enc}$, the 1D-convolutional layer has $256$ filters, a kernel size of $16$, and a stride factor of $8$. For masking network $\theta_{mask}$, the dimension $D$ of pre-processing is $256$. The segmentation splits the chunks with chunk size $S=250$ with $50\%$ overlap. For decoder $\theta_{dec}$, the transposed convolutional layer has the same kernel size and stride factor with the encoder.

We use dynamic mixing (DM) [28] and speed perturbation ($95\%$-$105\%$ randomly) for data augmentation. Adam optimizer [29] was used with learning rate of $1.5e^{-4}$. We also use automatic mixed precision [30] to speed up training.

In our experiments, Scale-Invariant Source-to-Noise Ratio improvement (SI-SNRi) and Signal-to-Distortion Ratio improvement (SDRi) [31] are used as the evaluating metrics. As listed in Table 2, we firstly explored different configuration of CA blocks number $N_{mask}$, IntraCA layers number $N_{intra}$, and InterCA layers number $N_{inter}$. The results indicated that our Tiny-Sepformer models achieve comparable separation performance with vanilla Sepformer on both WSJ0-2mix and WSJ0-3mix datasets, but have fewer model parameters. Furthermore, using the method of parameter sharing, the Tiny-SepformerS model could greatly reduce the model size, but with only a little performance degradation. All of the models are trained for $150$ epochs with batchsize $1$ on one NVIDIA V100 GPU card with $16$ GB memory. In particular, the $32$-layers Sepformer-32 is trained within $16$ GB GPU, instead of $32$ GB. We set a limit of training signal length $T$ to $64K$, to control the GPU memory consumption. For fair comparison, our Tiny-Sepformer-32 and Tiny-SepformerS-32 models used the same setting of this length limit.

As shown in Table 3, we also investigate different channel divisions of multi-head attention $D_{a}$ and separable convolution $D_{c}$ in IntraCA and InterCA network respectively. The best Tiny-Sepformer-32 and Tiny-SepformerS-32 models ($N_{mask}=N_{intra}=N_{inter}=4$) in Table 2 are used. The feed-forward dimension $D_{f}$ is $1024$, and the number of attention heads is $8$. The convolutional kernel size of IntraCA is $51$, and the kernel size of InterCA is $11$. The results demonstrated that large $D_{c}$ dimension ($D_{c}=192$) in IntraCA and $D_{a}$ dimension ($D_{a}=192$) in InterCA are better. The convolution helps the model to capture local information within each chunk, and the attention on the contrary models global context among all the chunks.

To further make the analysis of the function of convolution and attention paths, we plotted the weights of attention matrix in IntraCA and InterCA respectively. Best configuration in Table 3 (IntraCA $D_{a}=64$, InterCA $D_{a}=192$) of Tiny-Sepformer-32 is used in Figure 3. We found that the attentions of IntraCA are likely to gather together on the diagonal line (in Figure 3(a)). It means that large parts of attention could be replaced with local convolution in IntraCA. On the contrary, the attentions of InterCA are distributed globally among all the chunks (in Figure 3(b)). The huge difference between these two attention maps indicated that it is reasonable to assign more channel dimensions to $D_{c}$ for IntraCA and $D_{a}$ for InterCA.