Audio AI: isolating instruments from stereo music using Convolutional Neural Networks | by Ale Koretzky | Towards Data Science

This is the second article under the ‘Audio AI’ series I began back in March and it can be considered Part 2 after my first article on vocal isolation using CNNs. If you haven’t read that one yet, I highly recommend you start there!
As a quick recap, in that first article, I showed you that we can build a pretty-small-for-the-task Convolutional Neural Network (~300k parameters) to perform vocal isolation in real-time. We tricked this network into ‘thinking’ it was solving a simpler problem and eventually we got this kind of results:
Clip ID: 348095291
Delivery:application/vnd.vimeo.dash+json
Embed Size:692×270 
Separate AV:true
Dropped Frames:0 / 0 - 0
Playhead / Buffer:0 / 0 / 0
Bandwidth:0(0 Kbps0 Kbps)
Now, how do we go from here to reverse-engineering any professionally produced stereo mix into 4 instrument layers, more specifically, into the main 4 stems?
Clip ID: 348103015
Delivery:application/vnd.vimeo.dash+json
Embed Size:692×326 
Separate AV:true
Dropped Frames:0 / 0 - 0
Playhead / Buffer:0 / 0 / 0
Bandwidth:0(0 Kbps0 Kbps)
Before moving forward, a bit of context on Stems.
Stems are stereo recordings sourced from mixes of multiple individual tracks. For example, a drum stem will typically be a stereo audio file that sounds like all of the drum tracks mixed together. In most cases, additional processing such as equalization, compression, and time-based effects is included...[Izotope, on Stems vs Multitracks].
Where do stems come from? Well, stems have been used for decades in recording studios, as a way to control, process and manipulate entire groups of instruments, streamlining and simplifying the mixing, mastering and production process. Some will use stems, submixes, subgroups, or busses interchangeably. In addition to the workflow aspect,
stems have gained some popularity as a format on its own, that can expand the creative possibilities, particularly in remixing and DJing.
In this context, many players have been trying to develop the space, including Native Instruments with its STEM format.
As artists and labels become more comfortable and less threatened by a growing ecosystem that embraces collaboration and understands the creative and business potential of derivatives, stems have been experiencing some renewed attention.
The convention is to represent a mix by 4 stems:
I received great feedback for my first article’s journey structure so I’ll try to do something similar here. In this case, however, ‘A’ is now our trained, ready-to-go vocal isolation model and ‘Z’, the Stems isolation one.
This hypothesis should be broken down into 2 actually:
The same binary masking heuristics we used for isolating vocals from mixes, generalize well to other types of content. Or in other words, we expect comparable perceptual quality of isolation for all stems, using the same binary masking techniques.Regardless of whether 1) is true, the same CNN architecture we used for learning vocal binary masks will be able to learn the other stems’ binary masks with comparable accuracy.
The same binary masking heuristics we used for isolating vocals from mixes, generalize well to other types of content. Or in other words, we expect comparable perceptual quality of isolation for all stems, using the same binary masking techniques.
Regardless of whether 1) is true, the same CNN architecture we used for learning vocal binary masks will be able to learn the other stems’ binary masks with comparable accuracy.
If both 1) and 2) are true, we should be in a pretty good place.
Previously, I showed you that we can create binary masks that result in a perceptually decent reconstruction of the vocals when applied to the original mix. We then trained our CNN to learn how to estimate these binary masks, being that a much simpler task compared to estimating the actual vocal’s magnitude spectra.
In this context, our input-output representation and resulting training set looked like this:
Following the same process described in the first article for building our training set, we now have the same input as before but 4 sets of possible target outputs, one for each stem.
In this case, the youtube scraping approach didn’t give us much and we had to invest a lot of time and energy finding (and also synthesizing) perfectly aligned mix-stems pairs.
Let’s take the drums binary mask for instance. Now a value of ‘1’ indicates predominant presence of drums content at a given frequency bin — timeframe location, and a value of ‘0’ indicates predominant presence of something else at the given location. When talking about a frequency bin — timeframe location, remember we are literally referring to an individual element (pixel) in a binary mask, which corresponds to some value in the magnitude STFT domain.
Now that we have multiple sources, there are many different methods by which we can generate our binary masks. We actually ended up with 7 different variations, under 2 main principles:
a) Assume only one source can be present at a given frequency bin — timeframe location. In this scenario, you can think of all 4 sources competing for presence at every single location. Even though this is a huge over-simplification of reality, it works surprisingly well in some cases. (in terms of perception, the competition for presence is actually not that far from reality when thinking in terms of Auditory masking.)
b) Sources are allowed to share a frequency bin — timeframe location. Or in other words, there can be multiple sources present at any given location. This is obviously more in line with reality.
Since b) does not rely on the sources competing against each other, the binarization heuristics are based on a magnitude ratio criteria between each individual source and the original mix’s magnitude STFT at each frequency bin — timeframe location, and in some cases, conditioned on the neighbor regions in the time-frequency space.
Both principles have their pros and cons. The contents and demos in this article used a combination of proprietary methods under principle b). For a comprehensive overview and references on time-frequency masking, you can check this resource.
Now going back to our hypotheses, how do we validate hypothesis 1? Two approaches come to mind.
1) Using audio source separation metrics like SAR, SIR, SDR. These were introduced by the source separation academic community a few years back, with the goals of standardizing the criteria for assessing the quality of estimated sources recovered from a source separation process, as well as for the purpose of benchmarking. The museval toolbox implements these metrics.
2) A less formal (but not necessarily less effective when building a specific application) approach is to build small test sets and carry listening sessions comparing the original sources and the ones recovered through binary masking on the original mix. With this in place, you can establish some scoring based on the specific quality vectors you are interested in.
So what was the verdict for hypothesis 1? You’ll find out later...
Now that we have our Stems training set, how do we plug this input/output representation into some compatible architecture that’s going to learn the mix-to-stems task?
Well, once again, there’s more than one way to do this and I’ll show you a few approaches we went for during our design.
We know that our first model for vocal isolation worked well. Before thinking of reinventing the wheel, let’s first validate that the same CNN architecture we used for learning vocal binary masks behaves similarly with other stems. (hypothesis 2)
In essence,
We can train the exact same architecture we used for vocals, for isolating drums, bass and others, by simply changing our target binary masks
For instance, for the drums stem, we’ll be presenting X(mix) at the input and y(drums) at the output.
Once our 4 independent models have been trained, we can load the pre-trained weights into 4 identical model placeholders. From there we can connect these models in a way that they share a common input, while concatenating their output layers so that at every inference step, we obtain the binary masks for all 4 stems simultaneously. Think of these 4 models as becoming branches of a new, larger model (this is literally what happens here).
Looking at the output layer of this new model, you can see that our dimensionality is now 2052 (4 x 513). Because we know exactly how these model branches are wired, we can easily retrieve each stem’s binary mask by simple indexing.
Below, I show you the code to build the above architecture using Keras’ Functional API. I’ll leave it to you as homework to load the weights from a pre-trained model into new model placeholders. It should be relatively straightforward using the load_weights( ) / get_weights() / set_weights() functions in Keras.
from keras.layers import Dense, Dropout, Flatten, Conv2D, MaxPooling2D, Concatenatefrom keras.layers.advanced_activations import LeakyReLUfrom keras import Input, Model# this can be done in a for loop but I'll do it the 'redundant' way# for clarityinput_shape = Input(shape=(513, 25, 1))# BRANCH VOCALSx_1 = Conv2D(32, (3,3), padding=’same’, input_shape=(513, 25, 1))(input_shape)x_1 = LeakyReLU()(x_1)x_1 = Conv2D(16, (3,3), padding=’same’)(x_1)x_1 = LeakyReLU()(x_1)x_1 = MaxPooling2D(pool_size=(3,3))(x_1)x_1 = Dropout(0.25)(x_1)x_1 = Conv2D(64, (3,3), padding=’same’)(x_1)x_1 = LeakyReLU()(x_1)x_1 = Conv2D(16, (3,3), padding=’same’)(x_1)x_1 = LeakyReLU()(x_1)x_1 = MaxPooling2D(pool_size=(3,3))(x_1)x_1 = Dropout(0.25)(x_1)x_1 = Flatten()(x_1)x_1 = Dense(128)(x_1)x_1 = LeakyReLU()(x_1)x_1 = Dropout(0.5)(x_1)x_1 = Dense(513, activation='sigmoid')(x_1)# BRANCH DRUMSx_2 = Conv2D(32, (3,3), padding=’same’, input_shape=(513, 25, 1))(input_shape)x_2 = LeakyReLU()(x_2)x_2 = Conv2D(16, (3,3), padding=’same’)(x_2)x_2 = LeakyReLU()(x_2)x_2 = MaxPooling2D(pool_size=(3,3))(x_2)x_2 = Dropout(0.25)(x_2)x_2 = Conv2D(64, (3,3), padding=’same’)(x_2)x_2 = LeakyReLU()(x_2)x_2 = Conv2D(16, (3,3), padding=’same’)(x_2)x_2 = LeakyReLU()(x_2)x_2 = MaxPooling2D(pool_size=(3,3))(x_2)x_2 = Dropout(0.25)(x_2)x_2 = Flatten()(x_2)x_2 = Dense(128)(x_2)x_2 = LeakyReLU()(x_2)x_2 = Dropout(0.5)(x_2)x_2 = Dense(513, activation='sigmoid')(x_2)# BRANCH BASSx_3 = Conv2D(32, (3,3), padding=’same’, input_shape=(513, 25, 1))(input_shape)x_3 = LeakyReLU()(x_3)x_3 = Conv2D(16, (3,3), padding=’same’)(x_3)x_3 = LeakyReLU()(x_3)x_3 = MaxPooling2D(pool_size=(3,3))(x_3)x_3 = Dropout(0.25)(x_3)x_3 = Conv2D(64, (3,3), padding=’same’)(x_3)x_3 = LeakyReLU()(x_3)x_3 = Conv2D(16, (3,3), padding=’same’)(x_3)x_3 = LeakyReLU()(x_3)x_3 = MaxPooling2D(pool_size=(3,3))(x_3)x_3 = Dropout(0.25)(x_3)x_3 = Flatten()(x_3)x_3 = Dense(128)(x_3)x_3 = LeakyReLU()(x_3)x_3 = Dropout(0.5)(x_3)x_3 = Dense(513, activation='sigmoid')(x_3)# BRANCH OTHERSx_4 = Conv2D(32, (3,3), padding=’same’, input_shape=(513, 25, 1))(input_shape)x_4 = LeakyReLU()(x_4)x_4 = Conv2D(16, (3,3), padding=’same’)(x_4)x_4 = LeakyReLU()(x_4)x_4 = MaxPooling2D(pool_size=(3,3))(x_4)x_4 = Dropout(0.25)(x_4)x_4 = Conv2D(64, (3,3), padding=’same’)(x_4)x_4 = LeakyReLU()(x_4)x_4 = Conv2D(16, (3,3), padding=’same’)(x_4)x_4 = LeakyReLU()(x_4)x_4 = MaxPooling2D(pool_size=(3,3))(x_4)x_4 = Dropout(0.25)(x_4)x_4 = Flatten()(x_4)x_4 = Dense(128)(x_4)x_4 = LeakyReLU()(x_4)x_4 = Dropout(0.5)(x_4)x_4 = Dense(513, activation='sigmoid')(x_4)# concatenate output layersmerged_output = keras.layers.concatenate([x_1, x_2, x_3, x_4], axis=1)# connect to single/common input layermodel = Model(input_shape, merged_output)# define optimizers, loss, etc and compile model...
although the final model has close to 1.3M parameters, the learning process took place in architectures of ~300K parameters only
Draw your own conclusions but the points that I’d to make here are:
We can always try to solve a large problem by solving multiple smaller problems. (Not always the best thing to do in ML but sometimes it is)You definitely want to keep an eye at your number of parameters vs. the number of examples your model will be exposed to during training. Remember that in going from vocal isolation to stems deconstruction, all we did was generating alternative sets of target binary masks, but the number of training instances remained the same. With this design, we keep the same # training examples <> # parameters ratio that we had before.Because we trained each stem’s model in isolation, each model’s ability to learn its corresponding binary mask is not affected or influenced by the other models’ training process.
We can always try to solve a large problem by solving multiple smaller problems. (Not always the best thing to do in ML but sometimes it is)
You definitely want to keep an eye at your number of parameters vs. the number of examples your model will be exposed to during training. Remember that in going from vocal isolation to stems deconstruction, all we did was generating alternative sets of target binary masks, but the number of training instances remained the same. With this design, we keep the same # training examples <> # parameters ratio that we had before.
Because we trained each stem’s model in isolation, each model’s ability to learn its corresponding binary mask is not affected or influenced by the other models’ training process.
One thing you may be asking though... “Why not just keep those 4 models separate during inference?”
Well, it is a totally valid question. The answer is, for this particular application, we needed to perform the mix-to-stems deconstruction as part of a real-time pipeline connected to several other processes, including playback. Because of this, we wanted all stems to be generated at the exact same time to avoid buffering and synchronization issues. On top of that, we had to run this on a single GPU. Having all 4 models separate would have made things more challenging on both sides. If you’re building an API for example, and you don’t care about synchronization or real-time, you have a lot more flexibility.
Let’s say we want to stress-test our ability to generate all binary masks at once using the single stem architecture. What we can do is increase the output dimensionality to 2052, so that we can present all 4 binary masks at the output during training. This increase in the output dimensionality, however, will bring the total number of parameters up, to ~0.5M approximately. Not that this is necessarily an issue but, as someone who says “decouple your hypotheses!” 3 times a day, we should first see what happens when we train this model under comparable conditions.
There are many different ways in which we can try to reduce the number of parameters. We can, for example, decrease the dimensionality of the last FC layer before the output, from 256 to 77, as shown below:
During training, we present X(mix) at the input and a concatenation of all 4 stems’ binary masks at the output (2052 = 513 x 4). Some points to be made here:
The model will now have to figure out how to share its weights for fulfilling what each stem is demanding: its own binary mask. However, the learning process will now be global and not stem-specific.An important implication of 1) is that in some way, our stems are now competing for reconstruction quality. (not to be confused with our binarization principle a) where our stems compete for presence while building our binary masks). Even though we have a single vector at the output, the model’s ability to predict each stem’s binary mask will be implicitly affected by the other stems based on:
The model will now have to figure out how to share its weights for fulfilling what each stem is demanding: its own binary mask. However, the learning process will now be global and not stem-specific.
An important implication of 1) is that in some way, our stems are now competing for reconstruction quality. (not to be confused with our binarization principle a) where our stems compete for presence while building our binary masks). Even though we have a single vector at the output, the model’s ability to predict each stem’s binary mask will be implicitly affected by the other stems based on:
how consistent the input/output relationship between our mix and each stem is across our training set.
how complex the time-frequency patterns of each stem are in general, compared to other stems. From an MSE loss criteria, the network might ‘prioritize’ stems with easier masks to estimate, and/or converge to good reconstructions for certain stems before it does for others.
The 2 bullet points above are partially validated hypotheses, completely open to debate. They come from consistently observed behaviors during our experimentation.
In this case, the first 2 convolutional layers are stem-agnostic. The model has total discretion in how it will share those layers’ weights among all 4 stems. After the second convolutional layer, we create 4 branches of 2 convolutional layers each, one for each stem. Although we are not explicitly telling the model which branch corresponds to which stem, the model will be forced to optimize each branch to a specific stem based on the target binary masks at the output.
The reasoning behind this design was based on the known behavior of CNNs with structured data in the spatial domain, where the learning process derives into a hierarchical decomposition of features; from coarse and primitive to detailed and specific. Our hypothesis was, we can help the model learn better by persuading it to split the task between stem-agnostic filtering operations and stem-specific ones. By inspecting some of the feature maps out of the first convolutional stage, we observed no clear stem-specific patterns. However, as we moved deeper, most of the feature maps computed at the output of each branch revealed clear stem-specific patterns. This should not be a surprise, although it is very cool to see it :)
Like in the case of Model 2, the optimization is global and not stem-specific.
Perceptually, the results of this model are comparable to those from model 1. I’d say our hypothesis on the effectiveness of this particular design was partially validated but not conclusive.
For model 1, we merged independently trained models, remember?
What if we trained that same full architecture from scratch?
Not bad either. Some points worth mentioning here:
In this case, each branch’s task is stem-specific. For instance, the vocals branch will be learning how to isolate vocals and vocals only because of its target binary mask and because its weights are not shared or connected to any of the other stems’ branches.
However, each branch’s ability to learn its assigned task will be affected by the other branches... why? Just like in model 2 and 3, the optimization is global and not stem-specific. During training, the loss will be calculated for the full 2052-dimensional output vector, without any constraints to enforce ‘justice’ for all stems. This means that once again, the model might decide to prioritize certain stems over others based on their level of difficulty. Another way to read this is, some branches might converge to better solutions before others.
In the first article, I explained how our initial vocal isolation design ended up being some sort of regression-classification hybrid.
“We are asking the model to “classify pixels” at the output as vocal or non-vocal, although conceptually (and also in terms of the loss function used -MSE- ), the task is still a regression one”
By looking at our output representation, what prevents us from actually treating our problem as a classification one? At the end of the day, we need to choose between two discrete values at the output, 0 or 1. Therefore, we can transform our stems isolation problem into a multi-label classifier.
Think about the vocal isolation model for instance, with output dimensionality equal to 513. Think about what those values represent. We know they are all part of the same thing, right? Just points in a binary mask that later map into magnitude spectra, blah & blah... Well, another way to interpret those values is, they are labels or categories that represent the presence of content in different buckets.
An analogy might help crystallize this:
What changes do we need to make to transform our problem into a multi-label classifier?
Instead of using MSE as a loss function we’ll now be using binary cross entropy. The way you can think of this is, we now have 513 binary classifiers per stem, which will classify with a resolution of a single timeframe.
the final layer before the output should be followed by a sigmoid activation function. (This was the case already with our MSE-loss model but now it is a requirement)
Now, when we do an inference pass to estimate a middle frame’s binary mask, what we get is the probability of a given source being present per frequency-bin. In order to build the final binary mask, prior to the sliding window and concatenation mechanic described in the first article, we need to re-binarize these outputs by simple thresholding:
predicted_mask[predicted_mask >= 0.5] = 1predicted_mask[predicted_mask < 0.5] = 0
Treating our problem as a multi-label classifier results in some significant improvements in the quality of isolation compared to the MSE loss, regression — based model equivalents.
For this work, real-time operation was an important design constraint. This required us to spend a good amount of time doing software architecture design and implementation work, in addition to Machine Learning. Concretely, we had to ensure 2 things:
Making sure that the process of
... [reading audio, segmenting it, transforming it, predicting binary masks using our ML models, reconstructing stems back into time-domain and extracting additional features from the estimated stems including onsets, pitch, transcription, etc.. all this in Stereo] ...
was faster than the audio playback rate. In pursuing this, having a fast-inference ML model is obviously highly desired (therefore the reduced # of model parameters in our designs), although that is just one piece of the puzzle... hello multi-threading, sliding windows, buffering, synchronization, etc. To make this happen we built an optimized C/C++ pipeline on top of GStreamer, which was far from painless, but beautiful once we saw the light at the end of the tunnel. Our final solution ended up being ~2.5X faster than playback.
This is where things get real... In the first case, all we care about is performance, but you have no constraints in terms of access to your input data ahead of time. If your process is faster than your playback or delivery rate, you implement the required queueing mechanisms to save your data to retrieve it at a later time.
Now, what happens when your input is being streamed to you as it is being recorded and you need to perform demixing for a live music installation for example? Well, because you can no longer look ahead, if you need future context to predict the present, you’ll be introducing latency! What can we do to avoid (reduce) that latency then?
Going back to our input-output representations for training our models, one alternative is to ONLY rely on past context to learn to predict our binary masks, as shown on the left. That way, we pretty much limit latency to whatever recording/playback buffering stages our pipeline has.
The big question is: is the isolation quality affected when relying on the past only? Yes, significantly! And this should be no surprise, as music is time-series data. Any given timeframe will be more predictable if we know what happened before and after it. However, this still worked well enough for our purposes. [ Now you can picture me ranting on the importance of real Feature Engineering, like I did on my first article..]
Sorry to be this boring but they all have their pros and cons... Some models are more sensitive than others to certain types of content and mixing conditions. Others achieve better isolation at the expense of higher degradation in the reconstructed sources, etc. Overall, they all learned the task well enough with the exception of model 2, probably due to its limited number of parameters (~300k parameters) and high-dimensional output. In any case, my goal here was not to give you the holy grail of audio source separation (I’ll probably keep it to myself when I get there :)) but to provide you with some key insights and intuitions related to this particular problem, as well as to some considerations while designing artificial neural networks in general.
In terms of hypothesis 1, the final answer is, the binary masking method used for vocals generalizes relatively well for bass but not as well for drums and others. In most cases, the degradation you hear in the drums (especially in the highs/cymbals) is more related to the binary masking and not to the ML models. In the case of others, that same binary masking results in some degree of interference from other sources, especially vocals. You can hear this in the examples, which leaves in evidence that the ‘garbage in... garbage out...’ principle definitely applies in ML.
First of all, huge thanks to everyone who’s spent any time reading this article and/or the first one on vocal isolation. I received some great comments and feedback. Keep them coming!
Related to my first article, I found a couple of threads where some people pointed out other pieces of similar work, and I guess they were wondering whether I was aware of it. My answer is yes. But to clarify, my work, including the demos you’ve seen here, is actually from 2016... Those particular papers and presentations came after. Although my focus is on the Applied side of things, I’m well up to date (and very enthusiastic! ) about all the great academic work going on on this topic. With that said, even after 3 years, the kind of output quality you see here outperforms many of the current academic and commercial solutions out there.
If I had to solve this problem TODAY, however, I would probably explore other directions, maybe something recurrent in the time-domain, or a progressive GAN-based design, among other ideas. Binary /soft masking — based solutions will always hit a ceiling in the pursuit of perfect isolation.
Don’t hesitate in leaving questions. I’ll keep writing Audio AI stuff as time allows!
ps: once again, big shoutouts to Naveen Rajashekharappa and Karthiek Reddy Bokka for their contributions to this work.
disclaimer: intellectual property, designs and techniques described in this article have been previously disclosed in US Patents US10014002B2, US9842609B2, US10325580B2.