MULTI-SCALE TEMPORAL-FREQUENCY ATTENTION FOR MUSIC SOURCE SEPARATION

As shown in Fig.1(a), the neural network consists of an encoder, a separator and a decoder. It takes a discrete stereo signal with $N$ samples for each channel $\boldsymbol{y}\in\mathbb{R}^{N\times 2}$ as the input. The input signal is then transformed into a time-frequency domain representation $Y_{C}\in\mathbb{C}^{T\times F\times 2}$ via STFT, where $T$ is the number of frames and $F$ is the number of frequency bins of the complex spectrogram; 2 refers to the two-channel stereo input. To form the input of the neural network, the real and imaginary components are concatenated as $Y_{R}=[real(Y_{C});imag(Y_{C})]\in\mathbb{R}^{T\times F\times 4}$.

To reduce the computational cost for high-resolution input, following the subband mechanism proposed in [24], the fullband spectrogram is sliced into $K$ subbands to form the channel-wise subband signal $Y^{\prime}_{R}\in\mathbb{R}^{T\times(F/K)\times(4\times K)}$ served as the encoder input, where $K$ is set to 4 according to our preliminary results. The neural network estimates a channel-wise subband mask $M^{\prime}_{R}\in\mathbb{R}^{T\times(F/K)\times(4\times K)}$ and then reshaped to the fullband mask $M_{C}\in\mathbb{C}^{T\times F\times 2}$ for the stereo complex spectrogram. Finally, the target signal is estimated by multiplying the mixture signal $Y_{C}$ with the estimated mask $M_{C}$. The time-domain target signal $\hat{\boldsymbol{s}}$ can be obtained via iSTFT.

As shown in Fig.1(a), the encoder has three encoder blocks (EBs), each consists of a DenseNet block (detailed in Fig.1(b)) and a Conv2D block (detailed in Fig.1(c)). The DenseNet block consists of four Conv2D blocks (detailed in Fig.1(c)) with concatenation operations. The DenseNet block can learn explicit cross-layer interactions and reuses features computed in preceding layers, which yields efficient parameter utilization and suits for the MSS problems as discussed in [18]. The Conv2D block consists of a Conv2D layer, a Batch Normalization layer and an ELU activation layer.

The decoder consists of three decoder blocks (DBs) followed by a Conv2D block and a Dense layer. Each decoder block consists of a gated block [23] and a DenseNet block. The gated block consists of a Conv2DTanspose layer and two Conv2D blocks (detailed in Fig.1(c)) as shown in Fig.1(d), which learns a multiplicative mask for the feature from the encoder and suppress its undesired part. The structure of the Dense block in decoder is identical to the one used in encoder. After a Conv2D block, a Dense layer with tanh activation is utilized to generate the real and imaginary components of the complex ideal ratio mask (cIRM) with boundary of [-1, 1]. The cIRM is further extended to [-2, 2] by multiplying a expanding factor of 2 to increase the upper bound of the oracle SDR as discussed in [15]. The expanding factor is chosen to maximize the separation SDR according to our preliminary experiments. The detailed hyper-parameters of the encoder and decoder is listed in Table 1.

As shown in Fig.2(a), the temporal-frequency attention based separator consists of four residual attention (RA) blocks (detailed in Fig.2(b)). The input feature map of the RA block is $\mathcal{F}^{In}\in\mathbb{R}^{T^{\prime}\times F^{\prime}\times C}$ generated from encoder or previous RA block. $T^{\prime}$ is the time steps. $F^{\prime}$ is the feature dimension. $C$ is equal to 64. The input feature map is fed into two residual blocks. Each residual block consists of two Conv2D blocks with a kernel size of (3,3) and a stride of (1,1). The output of the residual block $\mathcal{F}^{Res}\in\mathbb{R}^{T^{\prime}\times F^{\prime}\times C}$ is then fed parallel into the temporal self-attention (TSA) and the frequency self-attention (FSA) blocks to capture the global dependencies along temporal and frequency dimensions, respectively. The outputs of the two self-attention blocks $\mathcal{F}^{Temp}\in\mathbb{R}^{T^{\prime}\times F^{\prime}\times C}$ and $\mathcal{F}^{Freq}\in\mathbb{R}^{T^{\prime}\times F^{\prime}\times C}$ are concatenated with $\mathcal{F}^{Res}$ to fed into a Conv2D block to generate the output of the RA block $\mathcal{F}^{RA}\in\mathbb{R}^{T^{\prime}\times F^{\prime}\times C}$, the kernel size and stride of the output Conv2D are all (1,1).

The TSA and FSA blocks share the same structure with different reshape operations as shown in Fig.2(c). The input $\mathcal{F}^{Res}\in\mathbb{R}^{T^{\prime}\times F^{\prime}\times C}$ is fed parallel into Conv2D blocks. The kernel size and stride of the Conv2D blocks are all (1,1), the channel number is reduced by half to $\frac{C}{2}$ for less computational complexity. The output feature maps of the Conv2D blocks in $\mathbb{R}^{T^{\prime}\times F^{\prime}\times\frac{C}{2}}$ are then reshaped to the $\mathcal{F}_{t}^{k}\in\mathbb{R}^{T^{\prime}\times(\frac{C}{2}\times F^{\prime})}$ for TSA or $\mathcal{F}_{f}^{k}\in\mathbb{R}^{F^{\prime}\times(\frac{C}{2}\times T^{\prime})}$ for FSA, respectively, where $k\in[{K,Q,V}]$. $K,Q,V$ indicates the key, query and value in the scaled dot-product self-attention [25]. For TSA, the self-attention is formulated as:

$$SA^{t}=Softmax(\mathcal{F}_{t}^{Q}\cdot(\mathcal{F}_{t}^{K})^{H}/\sqrt{\frac{C}{2}\times F^{{}^{\prime}}})\cdot\mathcal{F}_{t}^{V}$$

(1)

where $SA^{t}\in\mathbb{R}^{T^{\prime}\times(\frac{C}{2}\times F^{\prime})}$, $()^{H}$ denotes matrix transpose and $\cdot$ denotes matrix multiplication. For FSA, the self-attention is formulated as:

$$SA^{f}=Softmax(\mathcal{F}_{f}^{Q}\cdot(\mathcal{F}_{f}^{K})^{H}/\sqrt{\frac{C}{2}\times T^{\prime}})\cdot\mathcal{F}_{f}^{V}$$

(2)

where $SA^{f}\in\mathbb{R}^{F^{\prime}\times(\frac{C}{2}\times T^{\prime})}$. The $SA^{t}$ and $SA^{f}$ are further reshaped to $\mathbb{R}^{T^{\prime}\times F^{\prime}\times\frac{C}{2}}$ and then fed into a Conv2D block with channel number $C$, kernel size (1,1) and stride (1,1). The input $\mathcal{F}^{Res}$ is added to the output of the Conv2D block to get the final temporal self-attention $\mathcal{F}^{Temp}$ or frequency self-attention $\mathcal{F}^{Freq}$.

In section 2.3, the temporal attention is calculated based on all frequency bins. The frequency attention is calculated based on all input time steps. However, considering the different frequency ranges of various instruments and rapid changes of music content, using all frequency or temporal features might not be the optimal way for attention calculation. Therefore, multi-scale segment-wise attention is proposed to calculate attention based on different frequency or temporal ranges. More specifically, the input of TSA and FSA is first sliced into $P$ segments along the frequency and temporal dimension respectively. The attentions are then calculated for each segment individually and combined into the final output. The segment-wise TSA is formulated as:

$$\mathcal{F}^{Temp}=Concat(\mathcal{F}_{FP}^{Temp}(1),...,\mathcal{F}_{FP}^{Temp}(P))$$

(3)

$$\mathcal{F}_{FP}^{Temp}(i)=TSA(\mathcal{F}_{FP}^{Res}(i)))$$

(4)

where the $\{\mathcal{F}_{FP}^{Temp}(i),\mathcal{F}_{FP}^{Res}(i)\}\in\mathbb{R}^{T^{\prime}\times\frac{F^{\prime}}{P}\times C}$ is the i-th segment of $\mathcal{F}_{FP}^{Temp}$ and $\mathcal{F}_{FP}^{Res}$ respectively. The subscript $FP$ indicates that the $\mathcal{F}^{Temp}$ or $\mathcal{F}^{Res}$ is sliced into $P$ segments along frequency dimension. $TSA()$ is the temporal self-attention module as shown in Fig.2(c). The segment-wise FSA is formulated as:

$$\mathcal{F}^{Freq}=Concat(\mathcal{F}_{TP}^{Freq}(1),...,\mathcal{F}_{TP}^{Freq}(P))$$

(5)

$$\mathcal{F}_{TP}^{Freq}(i)=FSA(\mathcal{F}_{TP}^{Res}(i)))$$

(6)

where the $\{\mathcal{F}_{TP}^{Freq}(i),\mathcal{F}_{TP}^{Res}(i)\}\in\mathbb{R}^{\frac{T^{\prime}}{P}\times F^{\prime}\times C}$ is the i-th segment of $\mathcal{F}_{TP}^{Freq}$ and $\mathcal{F}_{TP}^{Res}$ respectively. The subscript $TP$ indicates that the $\mathcal{F}^{Freq}$ or $\mathcal{F}^{Res}$ is sliced into $P$ segments along temporal dimension.

The RA blocks with different value of $P$ can be combined together to form a multi-scale mechanism. Fig.2(a) shows the single-scale attention consists of four RA blocks with $P=1$. Fig.3 shows the separator with multi-scale attention, which consists of RA blocks with $P=1,2,4,8$. To obtain both small and large scale attention at the same layer, a two-branch structure with parallel RA blocks is introduced, where one branch contains RA blocks with increasing $P$ value while the other contains RA blocks with decreasing $P$ value.

A joint loss function with combining time domain and frequency domain losses is employed to train the network. The time domain loss is defined as mean absolute error (MAE) between target signal $\boldsymbol{s}$ and estimated signal $\hat{\boldsymbol{s}}$,

$$\mathcal{L}_{time}=||\boldsymbol{s}-\hat{\boldsymbol{s}}||_{1}$$

(7)

Where $\left||\cdot\right||_{1}$ denotes L1-Norm. The frequency domain loss is defined as MAE between target complex spectrum $\boldsymbol{S}$ and estimated complex spectrum $\hat{\boldsymbol{S}}$,

$$\mathcal{L}_{freq}=||real(\boldsymbol{S}-\hat{\boldsymbol{S}})||_{1}+||imag(\boldsymbol{S}-\hat{\boldsymbol{S}})||_{1}$$

(8)

The overall loss is defined as,

$$\mathcal{L}=\mathcal{L}_{time}+\alpha\cdot\mathcal{L}_{freq}$$

(9)

where $\alpha$ is set to 0.1 based on preliminary experiment results.

The proposed system is evaluated based on the MUSDB18 dataset [3], which consists of 150 songs with stereo format and 44.1kHz sampling rate. For each song, the final mixture signal is provided with its four audio stems, namely, vocals, bass, drums and other. We adopted the official split of 86, 14 and 50 songs for train, development, and evaluation respectively. The audio recordings are split into around 5.6 seconds (240 frames) segments with 2 seconds (86 frames) shift. The time domain segments are transformed to the time-frequency domain using an $8192$ samples STFT with $1024$ samples hop size. The complex spectrum $Y_{C}\in\mathbb{C}^{T\times F}$ of mixture clip is used as the input of system, where $T=240$ and $F=4096$. Data augmentation with channel swapping and remixing [26] is employed on the fly during model training.

For each audio source, we train a dedicated model individually. The Adam optimizer is employed. The initial learning rate is set to 0.001. It will decay by a factor of 0.8 when the validation loss does not decrease for 10 epochs. All the models are trained for 300 epochs.

We compare the proposed Multi-scale Temporal-Frequency Attention Network (MTFAttNet) with other existing start-of-the-art systems on MUSDB18 dataset in Table 2. The signal-to-distortion ratio (SDR) [27] computed by the museval toolbox[3] is used as evaluation metric. In the upper half of Table 2, the SDR results of the existing single domain systems including Spleeter[14], D3Net[12], Demucs [13] and ResUNetDecouple+ [15] are compared with the proposed MTFAttNet system. The $W$ and $S$ indicates the waveform domain and spectrogram domain respectively. In the bottom half of Table 2, we also listed the results of the top-ranked hybrid domain systems in the Music Demixing (MDX) challenge at ISMIR 2021 [28], namely, KUIELab-MDX-Net [29] and Hybrid Demucs[30]. $W+S$ indicates the method is working on hybrid domain.

As shown in Table 2, compared to the existing single domain methods, the proposed MTFAttNet method achieves significant improvement in separating vocals, drums and other stems. The bass SDR of the waveform domain method Demucs is higher than all the spectrogram domain methods, which might be caused by the limited frequency resolution of spectrum for bass. The overall SDR of the proposed system is 7.26 dB, which outperforms the best existing single domain method (ResUNetDecouple+ with SDR of 6.73 dB). Although the overall SDR is lower than the start-of-the-art hybrid domain systems, the proposed method still achieves the best vocal separation performance with the SDR of 9.51 dB, which outperforms the best existing hybrid domain method (KUIELab-MDX-Net with vocal SDR of 9.00 dB).

In this section, to better understand the benefit of the proposed MTFAttNet for MSS task, we further evaluate the systems by replacing the separator of MTFAttNet with different attention mechanisms in Table 3. Condition TFAttNet employed the temporal-frequency attention with single scale attention structure illustrated in Fig.2(a). FAttNet and TAttNet are conditions only using the frequency attention module (with temporal attention removed) and the temporal attention module (with frequency attention removed), respectively. Condition noAttNet removed both temporal and frequency attention modules from condition TFAttNet.

We first discuss the overall performances of different attention systems. As shown in Table 3, noAttNet achieves similar overall SDR (5.90 dB) to Spleeter (5.91 dB, listed in Table 2) indicating the effectiveness of proposed generic network structure without attention. With the frequency attention FAttNet and temporal attention TAttNet, the overall SDR is further improved to 6.65 dB and 6.79 dB respectively, which indicates the effectiveness of attention on capturing both temporal and frequency correlations. Combing both temporal and frequency attention, TFAttNet achieves the SDR of 7.09 dB. With multi-scale attention mechanism, MTFAttNet captures the spectrogram correlations more effectively and achieves the highest SDR of 7.26 dB.

The effects of attention mechanisms on different sources is varied for different type of stems. For the bass stems, we found that applying attention mechanisms introduced less improvements than other types of stems in Table 3. This is potentially caused by the limited performance of the spectrogram domain method for the bass stems as shown in Table 2. Based on noAttNet, the additional attention modules can not effectively capture the spectrogram correlations for bass with limited frequency resolution. For drums, FAttNet achieves the SDR of 6.44 dB while TAttNet achieves the SDR of 7.29 dB. The temporal attention outperforms the frequency attention for drums, which might be due to the fact that drums contain repeated beats in temporal domain. For vocals, compared to condition noAttNet, both temporal and frequency attention can significantly improve the performance. Combing both temporal and frequency attention can increase SDR to 9.23 dB. With multi-scale attention, the SDR is further increased to 9.51 dB. The success of attentions on vocals might benefit from the relatively long duration of pitch and the harmonic structure of vocals, which leads to the high correlations in both temporal and frequency domain. The results show such temporal and frequency correlations can be effectively modeled by the proposed MTFAttNet system.