DF-Conformer: Integrated architecture of Conv-TasNet and Conformer using linear complexity self-attention for speech enhancement

Let the $T$-sample time-domain observation $\bm{x}\in\mathbb{R}^{T}$ be a mixture of a target speech $\bm{s}$ and noise $\bm{n}$ as $\bm{x}=\bm{s}+\bm{n}$, where $\bm{n}$ is assumed to be environmental noise and does not include interference speech signals. The goal of SE is to recover $\bm{s}$ from $\bm{x}$.

In mask-based SE, a mask is estimated using a mask prediction network and applied to the representation of $\bm{x}$ encoded by an encoder, then the estimated signal $\bm{y}\in\mathbb{R}^{T}$ is re-synthesized using a decoder. The enhancement procedure can be written as

where $\operatorname{Enc}:\mathbb{R}^{T}\to\mathbb{R}^{N\times D_{e}}$ and $\operatorname{Dec}:\mathbb{R}^{N\times D_{e}}\to\mathbb{R}^{T}$ are the signal encoder and decoder, respectively, $D_{e}$ is the encoder output dimension, $\odot$ is the element-wise multiplication, and $\mathcal{M}:\mathbb{R}^{N\times D_{e}}\to[0,1]^{N\times D_{e}}$ is the mask prediction network. Early studies used the short-time-Fourier-transform (STFT) and the inverse-STFT (iSTFT) as encoder and decoder [9, 27], respectively. More recent studies use a trainable encoder/decoder [12] which are often called trainable “filterbanks” [28], e.g. in Asteroid  [13].

One of the main research topic in SE is the design of the network architecture of $\mathcal{M}$, because the performance and computational efficiency of SE are mainly affected by the structure of $\mathcal{M}$. Conv-TasNet [12] is a powerful model for speech separation and SE, and whose $\mathcal{M}$ consists of stacked 1D-DDC layers. TDCN++ [14, 15] is an extension of Conv-TasNet. The main difference of TDCN++ with Conv-TasNet is the use of instance norm instead of global layer norm and the addition of explicit scale parameters after each dense layer. The pseudo-code for $\mathcal{M}$ in the TDCN++ is shown in Algorithm 1. TDCN++ consists of $L$ stacked TDCN-blocks, and each TDCN-block mainly consists of two dense layers for frame-wise feature modeling and one 1D-DDC layer for sequence modeling. The dilation factor $d$ increases exponentially to ensure a sufficiently large temporal context window to take advantage of the long-range dependencies of the speech signal, and $L_{s}$ TDCN-blocks are repeated $R$ times where $L=RL_{s}$. The time complexity of TDCN++ is roughly proportional to $O(LN)$ when $N\gg D_{b}$, where $D_{b}$ is the input dimension of TDCN-blocks.

Conformer [21] is a derived model of Transformer [25] that was originally proposed for ASR [21] and later adopted in audio-related applications such as audio event detection [23, 24] and speech separation [29]. The structure of the Conformer is similar to the TDCN++, in that it consists of $L$ stacked Conformer-blocks [21]. Algorithm 2 shows the pseudo-code of a Conformer-block. By comparing Algorithm 1 and 2, we can see that the constituent layers of the Conformer-block and the TDCN-block are also similar; one Conformer-block mainly consists of several dense layers for frame-wise feature modeling, and one 1-D depthwise convolution layer and one MHSA-module for sequence modeling [21]. One of the main differences between the TDCN-block and the Conformer-block is the MHSA-module. Conformer enables global sequence modeling by using MHSA-modules instead of dilated depthwise convolution layers with local receptive fields.

Based on the successes of Conformer in speech-related tasks, we aim to replace the TDCN blocks in TDCN++ with Conformer-blocks. Unfortunately, the simple combination of trainable filterbanks and Conformer-blocks causes two critical problems. These problems are caused by the short window size of 2.5 ms and hop size of 1.25 ms used in trainable filterbanks for short-time analysis of the input signal.

Problem 1: The computational complexity. The computational cost of MHSA-module is quadratic in the number of frames $N$. In the original Conformer model [21], convolutional subsampling limits the size of $N$. For example, for a 1 second signal, $N$ is 25. In contrast, for TDCN++, the same signal would result in $N=500$.

Problem 2: The receptive field for sequence modeling is insufficient. The original Conformer has a hop-size of 40 ms, while the standard trainable filterbank has a hop-size of 1.25 ms. This means that the receptive field for depthwise convolution is 6.25 ms when using the default kernel size of 5, which may degrade the accuracy of the analysis of local changes in the signal.

One possible approach is to use the dual-path approach [16, 17, 18], which is equivalent to using sparse and block diagonal attention matrices corresponding to the inter- and intra-transformers, respectively. Alternatively, we use FAVOR+ attention introduced in Performer [26] which has linear computational complexity: $O(N)$. The novelty in our approach comes from using linear FAVOR+ attention to replace softmax-dot-product attention as well as performing local analysis with 1D-DDC to replace non-dilated convolutions in Conformer. Based on these two characteristics of the proposed method, we name our $\mathcal{M}$ as dilated-FAVOR-Conformer (DF-Conformer), and $L$-layer DF-Conformer is referred as DF-Conformer-$L$. The pseudo-code of DF-Conformer-$L$ is shown in Algorithm 3. The time complexity of DF-Conformer-$L$ is also roughly in proportion to $O(LN)$ when $N\gg D_{b}$.

Recently, many extended Transformer architectures have been proposed to make improvements around computational and memory efficiency [30, 31]. Performer [26] is one of them; it is an $O(N)$ Transformer architecture which uses FAVOR+. In self-attention, the query, $\mathbf{Q}$, key, $\mathbf{K}$, and value, $\mathbf{V}$ $\in\mathbb{R}^{N\times D}$ are combined as $\operatorname{sa}(\mathbf{Q},\mathbf{K},\mathbf{V})=\operatorname{softmax}(\mathbf{Q}\mathbf{K}^{\top})\mathbf{V}$. In FAVOR+, this is approximated as $\operatorname{sa}(\mathbf{Q},\mathbf{K},\mathbf{V})\approx\mathbf{D}^{-1}\phi(\mathbf{Q})\big{(}\phi(\mathbf{K})^{\top}\mathbf{V}\big{)}$, for a suitable feature map $\phi(\cdot)$ applied to the rows of each matrix, avoiding the quadratic term $\mathbf{Q}\mathbf{K}^{\top}$. Here $\mathbf{D}$ is a normalizing diagonal matrix with $\operatorname{diag}(\mathbf{D})=\phi(\mathbf{Q})(\phi(\mathbf{K})^{\top}\mathbf{1})$, and $\mathbf{1}\in\mathbb{R}^{N}$ an all ones vector. This approximation is made accurate in FAVOR+ using a random projection based non-negative valued $\phi(\cdot)$ of a suitable size [26]. To implement this idea, we replace the softmax-dot-product self-attention in Algorithm 2 with FAVOR+ self-attention. Hereafter, we refer to this new module as “MHSA-FAVOR-module”.

We strengthen the network’s temporal analysis capability by using 1D-DDC instead of the standard 1-D depthwise convolution used in the Conformer-blocks. As in TDCN++, we use an exponentially increasing dilation factor $d$. To implement this idea, DF-Conformer-block also takes $d$ as an argument, and it is passed to the 1D-DDC layer as the dilation parameter.

In a similar strategy as [22] and DF-Conformer, MHSA-FAVOR-module can also be incorporated into the TDCN-block. As an alternative network architecture, we insert $\bm{z}\leftarrow\bm{z}+\operatorname{MhsaFavorModule}(\bm{z})$ between line 9 and 10 of Algorithm 1, and refer to it as “Conv-Tasformer”.

Dataset: We used the same dataset used in the SE experiment of [15]. This dataset uses speech from LibriVox (librivox.org) and non-speech sounds from freesound.org . The duration of all samples were 3 sec, and sampling rate was 16 kHz. Training, validation, and test datasets consisted of 4,076,102 (3396.8 hours), 7,417 (6.2 hours), and 7,387 (6.2 hours) examples, respectively. We mixed speech and noise samples in the same manner of [32]. The minimum and maximum signal-to-noise ratio (SNR) of noisy input were $-40$ dB and $45$ dB, respectively, and the average extended short-time objective intelligibility measure (ESTOI) [33] score was 63.7%.

For the SE performance evaluation in Sec. 4.3, we compared DF-Conformer-$L$ and Conv-Tasformer with TDCN++ [14, 15] to confirm the superiority of the proposed models from its base model. In TDCN++, we used the same setting used in [32], namely, $L=32$, $L_{s}=8$, $D_{b}=256$ and $D_{c}=512$. In Conv-Tasformer, we used the same setting of TDCN++ except $L=16$ and $D_{r}=128$ to reduce the number of parameters.

For all models, $D_{e}=256$, and the window and hop sizes of trainable filterbanks were 2.5 ms and 1.25 ms, respectively. For STFT, the window and hop sizes were 30 ms and 10 ms, respectively, and fast-Fourier-transform size was 512. All models were trained for 500k steps on 128 Google TPUv3 cores with a global batch size of 512. We configured the Adam optimizer [35] with weight decay 1e-6, and learning rate schedule [25] of $D_{b}^{-0.5}\min(n\times 25000^{-1.5},n^{-0.5})$, where $n$ is a number of training steps. We clipped the gradient by global $\ell_{2}$ norm to 5.0. We stored a separate checkpoint of exponential-moving-averaged weights accumulated over training steps with decay rate 0.9999.

To confirm the effects of FAVOR+, we compared the real-time factor (RTF) of Conformer-4-STFT, Conformer-4, and F-Conformer-4 using 1 CPU. Figure 1 (a) shows the comparison results. In the case of Conformer-4-STFT, RTF does not increase significantly because $N$ was $100/\mbox{sec}$ in our STFT setting and it is still feasible with $O(N^{2})$ MHSA-module. Whereas RTF of Conformer-4 increases linearly as $N$ was $500/\mbox{sec}$ in our trainable filterbank setting and MHSA-module. Since the time-complexity of FAVOR+ is in proportion to $O(N)$, F-Conformer-4 has solved this problem.

We also compared these methods using two objective metrics; scale-invariant SNR improvement (SI-SNRi) [36] and the ESTOI. Table 1 shows the results. By comparing Conformer-4-STFT and Conformer-4, the use of a trainable filterbak achieved higher scores than STFT as similar to previous studies [14]. When using the small-size model, the SI-SNRi score of F-Conformer-4 was almost the same as those on the Conformer-4-STFT. Meanwhile, with the 8.75M models, SI-SNRi of F-Conformer-8 was 1.2 dB higher than that of Conformer-8-STFT, and ESTOI scores of those were almost comparable. These results suggest that the use of FAVOR+ can achieve high time-domain SE performance with a larger model while avoiding the increase in computational complexity.

We compared DF-Conformer-8, TDCN++, and Conv-Tasformer using SI-SNRi, ESTOI, and RTF. From the comparison results shown in Table 2, DF-Conformer-8 and Conv-Tasformer achieved comparable scores, and these scores were higher than that of TDCN++. Also, by comparing DF-Conformer-8 and F-Conformer-8 in Table 1, the use of 1D-DDC significantly improved the scores while avoiding to increase RTF. These results suggest that the use of both 1D-DDC and FAVOR+ is effective in SE. We also compared RTF of these methods as shown in Fig. 1 (b). RTFs of DF-Conformer-8 and TDCN++ were comparable, whereas that of Conv-Tasformer was larger than others due to additional MHSA-FAVOR-block. Therefore, when inserting FAVOR+ in TDCN-block as Conv-Tasformer, it will be necessary to devise the position and number of MHSA-FAVOR-module in order to improve the computational efficiency.

We also compared the iterative extension of these models [14]. Using iterative model improved the scores of all methods, and the results tended to be similar to the non-iterative models. Furthermore, we evaluated a larger model as iDF-Conformer-12 with $L=12$, $D_{b}=256$, and the number of attention heads were 8. The size of model was determined so that RTF becomes comparable with iConv-Tasformer. As we can see the results, the scores clearly improved using a large model, thus DF-Conformer would be able to scale the performance according to the model size.

We finally point out three characteristics in DF-Conformer’s attention matrices. First, none of all attention matrices has a local structure that focuses only on nearby time-frames. Secondly, most attention matrices in earlier layers referred to low SNR time-frames to capture the noise characteristics (e.g. Fig. 2 middle-left), or referred to time-frames with similar spectral structures (e.g. Fig. 2 middle-right). Thirdly, some attention matrices of deep layers resemble a sum of a nearly-diagonal matrix and a block matrix (e.g. Fig. 2 bottom). This results suggest that the earlier layers roughly analyze the speech and noise from the entire utterance, and later layers refine the mask based on the local structure.