Few-Shot Musical Source Separation

We use a U-Net architecture very similar to that described in [11] as the baseline. The model is conditioned directly on instrument classes (via one-hot encoding) and can be viewed as a standard supervised learning baseline. Here, we briefly describe our variant of the model. For the U-Net architecture, each of the 6 encoder layers consists of $n_{i}$ 5x5 2-D convolutions with a 2x2 stride, followed by batch normalization and a Leaky ReLU activation, where $n_{0}=16$ and $n_{i}=2n_{i-1}$ for each subsequent layer $i$. The decoder mirrors the encoder by replacing 2-D convolutions with 2-D transpose convolutions. Additional skip connections are added between layers at the same hierarchical level in the encoder and decoder.

A one-hot vector that directly encodes the target instrument class is used as the conditioning vector $z_{\textit{class}}$. A FiLM layer is inserted after the U-Net bottleneck layer to transform the bottleneck feature $x$ via $FiLM(x)=\gamma(z_{\textit{class}})\cdot x+\beta(z_{\textit{class}})$, where $\gamma$ and $\beta$ are learned parameters that scale and shift $x$ based on the conditioning information $z_{\textit{class}}$. In our work, $z_{\textit{class}}$ is 18-dimensional to support a wide variety of instrument classes. The y-axis in Figure 2 shows 14 of these instrument classes, and the remaining 4 are mallets, pipe organ, bagpipes, and whistling.

The model takes three seconds of mono audio with 22.05 kHz sample rate as input and first computes a spectrogram with a hop size of 256 and a FFT size of 1024. The spectrogram is log-compressed, and input into the U-Net encoder. The output of the decoder is treated as a two-channel complex mask, where a sigmoid activation is applied on the magnitude as in [22]. The mask is applied to the compressed spectrogram via complex multiplication, followed by decompression. Finally, the Inverse STFT is computed to retrieve the separated audio signal. The loss is computed as the sum of the SDR Loss [23] in the time domain and mean absolute error on the magnitudes in the time-frequency domain.

In our proposed few-shot MSS paradigm, the external information is based on few audio examples of the target instrument instead of its class label. We train a few-shot conditioning encoder jointly with the U-Net to generate a learned conditioning vector $z_{\textit{few-shot}}$ based on few target examples.

The conditioning encoder consists of four convolution blocks, each of which has a 64-filter 3x3 convolution, a batch normalization layer, a ReLU activation layer, and a 2x2 max-pooling layer. To allow our model to handle variable-length input, we apply max-pooling along the time dimension to the output of the convolution blocks. Finally, we flatten the feature map to output an embedding of size 512. When multiple audio examples are used as conditioning (few-shot instead of one-shot), each example is used to generate a conditioning vector, and the example-wise vectors are averaged into a single conditioning vector. We keep the conditioned U-Net architecture, FiLM conditioning mechanism, and loss terms the same as in the baseline model to systematically compare the instrument class conditioning and few-shot conditioning paradigms.

In order to support a wide variety of instrument classes, we train the models on a combination of several multitrack datasets consisting of solo-instrument stems that are linearly mixed to create full song recordings. The first dataset, $D_{\textit{real}}$, is a private dataset containing 313 real-world multi-tracks, with an instrument bias towards pop/rock music meaning that vocals, drums, bass, guitars, and synthesizers dominate the distribution. The second, $D_{\textit{synth}}$, is the public Slakh2100-redux [26] dataset, containing 1710 multi-tracks synthesized from MIDI with high-quality synthesizers. This dataset contains no vocals, and by its nature only contains synthetic versions of each instrument. $D_{\textit{solo}}$ is a private dataset of 627 real solo instrument recordings, roughly evenly distributed across the 18 instrument classes. These recordings are randomly combined to generate multi-tracks with up to 5 instruments.

To generate a training example, we randomly sample a dataset $D$ from the set $\{D_{\textit{real}},D_{\textit{synth}},D_{\textit{solo}}\}$, a multi-track $M$ from $D$, a chunk from $M$ as the input mixture, and finally an instrument track from the mixture as the target output. To train the few-shot models, we additionally sample $n$ random chunks from the original track of the target instrument as the conditioning examples. Note that we make sure the random chunks do not overlap with the target output. We experiment with $n$ between one and five (1-shot to 5-shot) to examine how this choice correlates with model performance. All models are trained using a batch size of 16 and a learning rate of 0.001 with the Adam optimizer and early stopping.

We evaluate the trained MSS models on the 50 test tracks in the MUSDB18 dataset. For each track, we have its mixture and isolated tracks of vocals, drums, bass, and other in stereo format with a 44.1k Hz sample rate. An other track includes all instruments in a mixture other than vocals, drums, and bass. It can vary greatly over time with different instruments present at different times, making conditioning using random examples potentially incoherent. Therefore, we do not evaluate on separating other in this work. Note that we did not include the training set of MUSDB18 into our training data. The training and testing data are therefore from different distributions, resulting in a more challenging setup compared to most of the previous works on MSS.

To evaluate a test track, we downsample the audio to 22 kHz and divide the mixture into 3-second input segments with 50% overlap. We sample $n$ random chunks from the target instrument track as the conditioning examples for the few-shot models, and run inference on each 3-second chunk. We reconstruct the final signal using overlap-add, and upsample to 44.1 kHz. We compute the signal-to-distortion ratio (SDR) [27] per track using the museval package as our evaluation metric, and average the SDR across tracks. We repeat this process 10 times for each track with different random conditioning examples when evaluating few-shot models, reporting the mean and standard deviation of means over each of the 10 test iterations.

While MUSDB18 is the most common MSS benchmarking dataset, it only contains isolated tracks for three instrument classes. Since our baseline model supports up to 18 instruments, we additionally evaluate the baseline and the best performing few-shot model on the MedleyDB dataset. We filter out multi-tracks with bleed, map all instruments to the predefined 18-instrument vocabulary, and generate sub-mixes in which we keep only one stem per instrument class. The resulting MedleyDB evaluation set contains 90 tracks with 14 instrument classes.

In our basic setup of training and evaluating the few-shot model, the selected conditioning examples have two main constraints: (1) they are drawn from the target recording, (2) they are single-sourced audio examples that only contain the target instrument. This is a relatively ideal scenario when considering real-world applications, where isolated examples from the target recording may not be available. To see how a few-shot model generalizes to more challenging and practical scenarios, we perform a set of experiments where we evaluate the model with one constraint relaxed at a time.

To relax the first constraint, we draw $n$ conditioning examples from $n$ different recordings of the same instrument class within the test data. To relax the second constraint, we allow some conditioning examples to contain one additional non-target instrument, where each example has a 0.5 probability of being multi-sourced. To get a multi-source example from a multitrack, we randomly choose a non-target instrument and mix it with the target before we sample an example chunk. Note that we make sure only one consistent instrument, the target, is present across all conditioning examples.

Besides providing target instrument examples to specify the desired output, we further explore if specifying unwanted instruments through the same conditioning mechanism can help improve the separation. To do so, we extend our basic few-shot model to leverage additional negative conditioning examples. During both training and evaluation, we draw $n$ additional negative examples, each of which contains a random non-target instrument that also exists in the input mixture. All conditioning examples are first embedded via the same few-shot conditioning encoder. Then we average the example-wise conditioning vectors over positive and negative examples separately. We concatenate the aggregated positive and negative conditioning vectors along the feature dimension, and convert the concatenated vector to a final conditioning vector via a fully-connected layer with a ReLU activation. The sizes of the resulting conditioning vector and the conditioned U-Net model are the same as in our basic setup.

In Table 1, we show the performance of the instrument class conditioned baseline and few-shot conditioned models with different numbers of conditioning examples $n$ (the number of shots).

First, the 1-shot model, conditioned on only one audio example, does not show a significant advantage over the baseline. It achieves higher SDR on vocals, but performs worse on bass. While we can provide more direct and specific information about the target via audio examples for conditioning, one random example may not be sufficient to capture the variations of the target within the entire song. On the other hand, as $n$ increases, the few-shot model begins to outperform the baseline across all instruments. This shows the advantage of extending existing query-based or one-shot MSS to few-shot MSS. We can now provide more conditioning examples to better capture the instrument variations in pitch, timbre, and playing technique, to achieve better separation. Next, we see a general trend that as $n$ increases, the performance of the few-shot model initially increases and then stays relatively stable, while the standard deviation continues to decrease. Since we aggregate example-wise conditioning vectors by taking an average, we can expect these diminishing returns from including more examples once enough representative examples are present. On the other hand, conditioning on more audio examples can mitigate the randomness from sampling and leads to more stable results. Therefore, in the remainder of our experiments, we focus on a 5-shot model for more robust performance.

In addition to MUSDB18, we evaluate the baseline and 5-shot models on MedleyDB, which contains many more instrument classes. Figure 2 (left) first shows that the 5-shot model outperforms the baseline on 13 instruments out of the total of 14 with an overall SDR of 1.96 dB (0.56 dB for the baseline). Next, we sort the instruments by their occurrences in the training data based on the distribution shown in Figure 2 (right). We see that the 5-shot model can not only effectively separate instruments that are relatively rare in our training data, but also generally outperforms the baseline by larger margins on these rare classes, for example, Keyboards, Woodwinds, and Accordion. This suggests that the advantage of few-shot MSS is more significant on rare instruments. From these results, we see that few-shot MSS is applicable to many instruments across different datasets. In the remainder of our experiments, for brevity we focus on the MUSDB18 dataset.

One of the advantages of the few-shot MSS model is that it is able to handle new, unseen instruments at inference time without being limited by the training vocabulary. However, our training data has a large instrument vocabulary while MUSDB18 only contains three common instruments. To quantitatively evaluate how well models can separate unseen instruments, we remove all vocals, drums, and bass target examples from our training data and retrain the baseline and few-shot models with the updated training set. This ensures that the model does not see instruments from the test set during training.

The results in Table 2 first show that the baseline model conditioned on instrument classes, unsurprisingly, does not perform well on unseen instruments as it is effectively untrained for these classes. On the other hand, both few-shot models with $n\in\{1,5\}$ outperform the baseline by large margins across all instruments (except the 1-shot model for drums). Note that the performance on drums is significantly lower than the other two instruments. We conjecture that the high variance of drum sounds makes it harder for the model to separate based on just one or few examples at inference time. In addition, if we compare Table 2 to Table 1, first, we see the drops in SDR resulting from not seeing the testing classes during training. Next, we see that the 5-shot model outperforms both the baseline and the 1-shot model more significantly on unseen instruments. This aligns with our observations in Section 4.1 that the advantage of few-shot MSS is more significant on rare and even unseen instruments.

In the previous experiments, we provided few-shot models with single-sourced conditioning examples drawn from the target recording. To see if few-shot models are able to generalize to more practical scenarios, we evaluate a 5-shot model with relaxed constraints on conditioning examples as discussed in Section 3.4.

Table 3 first shows the performance of the basic 5-shot model where (as before) we apply both constraints to the conditioning examples. Then, if we draw conditioning examples from different tracks instead of the target track, SDR drops between $\approx$0.4 to 1.1 dB across instruments. Even with this drop, the model still achieves comparable and in some cases better SDRs compared to the baseline model (shown in Table 1). On the other hand, if we provide multi-sourced conditioning examples with additional non-target instruments instead of single-sourced ones, the model performance drops significantly. This is not surprising given that the model is trained only on isolated audio. To see if changing the training conditions would make this scenario possible, we retrain the 5-shot model with multi-sourced conditioning examples. The retrained model (shown as MS in Table 3), while still performing worse than both the basic 5-shot model and the baseline, achieves about 0.6 to 2.8 dB SDR improvement by matching the training and testing scenarios. These results indicate that the proposed few-shot MSS model can adapt to a wide variety of applications where the conditioning examples can be drawn from different recordings. The model has a harder time generalizing to multi-sourced conditioning examples, but if that is the desired scenario, we can match the training and inference objectives to achieve better results.

Next, we experiment with including five additional negative conditioning examples, each containing a random non-target instrument, as discussed in Section 3.5. The results at the bottom of Table 3 show that providing additional information about what not to separate during training and inference helps to further improve the separation compared to the basic 5-shot model conditioned on positive examples only. This setup could be particularly useful when labeling non-target instruments is much easier than labeling the target one.