Consistent and Relevant: Rethink the query embedding in
general sound separation

As shown in Fig.1, our CaRE-SEP is composed of shared encoder, connector module and decoder module, where the encoder and connector module are responsible for query embedding, and all of them are trained for the audio source separation.

In Fig.1, we regard query embedding network as a part of separation model and use a shared encoder module to simultaneously obtain query embedding and separation feature, which is different from [4] that uses a frozen extra classification model (ST-SED) to obtain query embedding. When the query embedding module and separation structure are similar, the distribution of generated query embedding and separation feature will be more consistent, leading to optimal performance. Furthermore, we initialize the shared encoder of CaRE-SEP with a pre-trained classification model [15] and perform gradient propagation for query embedding. As a result, the query embedding is optimized to be more relevant to separation task. Next, we will introduce the training process depicted in Fig.1.

We define two audio sources $a_{1},a_{2}$ and mix them $a=a_{1}+a_{2}$ with energy normalization. To balance classification and separation performance, we do not use mel-spectrograms as input which perform worse in separation. Instead, we use Mel scale-based band-split [16] to split the spectrogram into a series of subband spectrograms. The mixture audio $a$ is converted into the spectrogram by Short-time Fourier Transform (STFT) and then split into subband spectrograms. Besides, we use anchor audios to obtain fixed-dimensional query embeddings $e_{1},e_{2}$, which are incorporated into the CaRE-SEP to specify which audio needs to be separated. This process is illustrated by the dashed line in Fig.1. Subsequently, we send two training triplets $(a,a_{1},e_{1}),(a,a_{2},e_{2})$ into the CaRE-SEP model S, respectively. Our separation objective is as follows:

Simultaneously, the linear layer (embedding layer in Fig.1) is applied on the query embedding $e_{j}$ to unify the feature dimension, and then the queries are added into the audio feature maps of each layer block in the CaRE-SEP model, helping to specify the audio that needs to be separated. As a result, the model will also learn the relationship between the query embedding and mixture, adjusting its weights to adapt to the separation feature. Finally, the output spectrogram is converted into the waveform by inverse STFT (iSTFT). We apply the Mean Absolute Error (MAE) and Source-to-Distortion Ratio (SDR) to compute the loss between separate waveforms and the target source.

As illustrated in the left part of Fig.2, our shared encoder mainly consists of Swin Transformer [17] and Patch Merge [15, 4]. The Patch-Embed CNN splits the input features into different patch tokens inside each window, which more effectively capture the relationship among frequency bins of the same time frame. The transformed patch tokens are sent into three groups of Swin Transformer blocks [17] and Patch-Merge layers to generate hierarchical representations. Within each group, we adopt the more efficient swin transformer block with a shifted window attention to learn feature representation. Specifically, the patch-merge layer is applied to decrease the sequence size at the end of each group.

The connector module contains a swin transformer, token-semantic module [4], and average pooling layer. Hierarchical representations of the swin transformer have the ability to address multiple tasks. Firstly, hierarchical representations can be input into the decoder network to separate isolated audio sources from mixed audio. Besides, they can be processed through the average pooling layer to generate query embedding for separation. Additionally, the token-semantic CNN integrates all frequency bins and maps the feature into the class probabilities. The vectors produced by the token-semantic CNN can accomplish audio tagging tasks by computing binary cross-entropy loss with ground truth labels. So we can simultaneously train tagging and separation tasks in a single model.

Taking inspiration from U-Net [18], the decoder of our CaRE-SEP is a structure symmetrical to our encoder, which can better accomplish separation tasks. It is composed of three Swin Transformer blocks and three patch expand layers. Different from the patch merge layer, the patch expand layer is devised for up-sampling, which transforms adjacent feature maps into larger feature maps with 2x up-sampling of resolution. Finally, the patch up layer performs 4x up-sampling to restore the feature maps to the original input resolution. The skip connection combines extracted representations of the decoder and multi-scale features from the encoders together, which helps to complement the loss of spatial information resulting from down-sampling.

AudioSet. Our CaRE-SEP model is trained on AudioSet [19], a large-scale dataset of over two million 10-second audio samples labeled with 527 sound event categories. In the following experiments, we use the full-train set of AudioSet (2M samples) to train our model and evaluate its performance on the evaluation set (22K samples).

In this experiment, we train our CaRE-SEP model on AudioSet and evaluate the performance, with the results presented in Table 1. The hidden dimensions of the separation model have different sizes, including 527-d, 768-d, and 2048-d.

For the audio classification task, when we only train audio tagging with CaRE-CLS, the performance is almost close to the best result. The difference between them is that 768d ST-SED uses mel-spectrogram while we use subband spectrogram as input. When we perform joint training in CaRE-CLS+CEP, the classification performance is slightly worse, the separation performance is better than the AudioSet baseline but not as good as our CaRE-SEP, indicating that there are some conflicts between the classification feature and the separation feature. We will further explain this in Section 4.2

For the separation task, we could clearly figure out that our CaRE-SEP models trained at 768-d and 2048-d embeddings achieve optimal performance. Compared to the baseline (2048-d ST-SED-SEP), our best mixture-SDR increases by more than 3 dB, and the improvement of clean-SDR is quite significant, demonstrating that our model can maintain the audio features very well. All of these results demonstrate the effectiveness of our proposed model.

It is worth noting that our model exhibits a similar performance in both the 768-d and 2048-d dimensions. A possible reason is that the larger embedding space may have some redundancy, and the 768-d query embedding space can already help the model better capture the audio features and provide more discriminative embeddings for more accurate separation.

In this section, we aim to analyze the effectiveness of our proposed model from the perspective of feature distribution. We randomly sample 70 audio samples from each of 2 categories in the AudioSet evaluation set and extract separation feature and query embedding from different models. As shown in Fig.3, we use t-SNE [26] to visualize the distribution of embeddings, with different colors indicating individual categories. In the 2048-d ST-SED-SEP (baseline) and our CaRE-SEP, we extract output of the final block in the connector module as the separation feature.

Fig.3(a) and Fig.3(c) indicate that separation feature can slightly discriminate categories, with a small inter-class difference. However, the query embedding shows more apparent class discriminability in Fig.3(b), which validates the conjecture that a mismatch exists between the query embedding and separation feature in the baseline. Fig.3(d) is almost consistent with Fig.3(c), which further suggests that our model can generate consistent and relevant query embedding with separation feature. In addition, we also calculate the KL divergence between the embeddings of two categories. The four KL divergences in Fig.3 are (a)4.34, (b)10.37, (c)5.08, (d)5.77, respectively. This metric analysis indicates that the distribution between separation feature and query embedding exists mismatches in baseline. While in our CaRE-SEP, their distribution is closer.

Moreover, Fig.3(b) is obtained from the extra classification model. Fig.3(b) and Fig.3(a) can indicate the classification features are inconsistent with the separation feature. This can also explain the reason for the slight decrease in classification performance when jointly training classification and separation in Section 4.1.

As demonstrated in Table 2, we investigate the impact of query embedding on the performance of separation models from three perspectives. $Grad$ refers to the gradient backpropagation during the training (not freezing). $Init$ is to initialize the shared encoder with a pre-trained classification model. $Shared\ encoder$ means that the separation encoder can simultaneously generate query embedding.

We validate the effectiveness of the proposed swin-unet architecture, shared encoder, classification initialization, and gradient backpropagation sequentially. First, the proposed swin-unet architecture is tested by comparing models A and G, each of which has an extra classification network [4]. Model G gets comparable results with the baseline, demonstrating the effectiveness of our proposed model structure. Second, the shared encoder is tested by comparing models B and C. Model B trains the query embedding with gradients in the extra network and performs worse than model C, indicating the importance of encoder sharing. Third, the initialization of the shared encoder is evaluated by comparing models D and F. Model D uses a pre-trained classification model to initialize the shared encoder, which improves the performance by nearly 1 dB compared with model F. This confirms that the classification embedding can help audio source separation. Fourth, we train the query embedding with gradients in the shared encoder to validate the gradient backpropagation. Model E demonstrates a significant increase in mixture SDR when compared to models F and D, indicating that gradient backpropagation is crucial for query embedding. Moreover, we test model F with a shared encoder to directly generate query embedding, which results in the worst mixture SDR. This suggests that $Grad$ and $Init$ are indispensable for the shared encoder.

To verify generalization ability, we further evaluate the zero-shot performance on ESC-50. Table 3 shows our CaRE performs better than baseline, they can both demonstrate the effectiveness of consistent and relevant embeddings. The performance of joint training is better than CaRE-SEP, which is different from the results on AudioSet. This can be explained by that these two datasets are inconsistent. The models trained on AudioSet train set perform well in its evaluation set. While directly testing the zero-shot performance on ESC-50, the performance of CLS+SEP is better, indicating labels play an important role in zero-shot generalization.