TFSWA-UNet: Temporal-Frequency and Shifted Window Attention Based UNet For Music Source Separation

The performance of music source separation (MSS) has been greatly improved in recent years due to the rapid development of various neural network architectures. Spectrogram features are widely used in MSS tasks which consists of both time and frequency information. However, the time and frequency correlations and the local patterns of spectrogram have not been fully explored. In this paper, we propose a novel U-Net architecture named TFSWA-UNet in which a temporal-frequency and shifted window attention (TFSWA) based block is designed as bottleneck block. In the proposed TFSWA block, time sequence attention (TSA) block and frequency sequence attention (FSA) block are used to capture the global correlation of music spectrogram features within time and frequency sequences respectively. To further capture the local correlations of spectrogram features, shifted window attention based Swin transformer is also introduced into the TFSWA module, which computes self-attention within local non-overlapping windows and capture correlations from both temporal and frequency dimensions. Experiment results on MUSDB18 dataset indicate that the proposed TFSWA-UNet model outperforms the state-of-the-art MSS methods with SDR of 9.16dB on vocals.