Transformer Neural Net makes music! (JukeboxAI)

take me to church I'll worship like a dog at the shrine of your lies I'll tell you my sins so you can sharpen your knife offer me my deathless death good God let me give you my life no masters or kings when the ritual begins there is no sweeter innocence than our gentle sin in the madness and soil of that sad earthly scene only then I am human only then I am clean there's this AI that can generate music in the style of famous singers it can make it rapping Bruno Mars a rock-and-roll Katy Perry and so much more but how does it do this we're going to try to reconstruct this model intuitively and get into more details as we go in doesn't matter how much you know about AI you should just be able to walk away with some knowledge of how this works at any level of detail so let's get started so in the first pass we'll build this intuitively we have this AI that takes in some lyrics genre and an artist and it generates a song and it does so in fixed sized chunks the AI generates a chunk it can then take this chunk and use it to generate the next chunk of audio and it repeats this until we get the entire song this type of AI model takes in an audio sequence and generates the next audio sequence and hence it is a sequence to sequence model in deep learning literature transformer neural networks are the best at dealing with sequence data at least for now so we can replace the AI model with a transformer neural network pretty cool but we run into a problem when building this out this raw audio waveform it's huge you've seen songs on CDs have a sampling rate like forty four point one kilo Hertz that means that we use 44100 numbers just to represent one second of audio it's far too big for our model to handle so we need to compress the waveform while retaining the key aspects of the music and this compression is done through a type of neural network architecture called an autoencoder this autoencoder network can take in a raw waveform and learn to compress it and it can also learn to take a compressed audio and learn to decompress it back to the original waveform the jukebox uses the auto encoder and transformer together and we build this jukebox in two phases training and generation during the training phase we train the auto encoder to compress and decompress audio and then we train the transformer to take in some information about the song to generate and train it to generate a compressed vector representation of the song one chunk at a time during the generation phase though we use both of them together pass in the lyrics genre and artists to the transformer and have it generate the compressed vector 1 audio chunk at a time and we can pass each of these compressed vectors through the auto encoder to get back the raw audio waveforms stitch them together and we have the generated song if we build our jukebox out in this way it might sound something like this hey jukebox how the anonymous celebrity rap to lose yourself by Eminem music to my ears it sounds like this because we built this only based on intuition but there are some key details and modifications we need to make before this is usable and now let's take a look at these in pass to building the real architecture we're using an auto encoder to compress and decompress a waveform this is split between two parts of an auto encoder its encoder for compression and it's decoder for decompression during training the normal auto encoder seeks to minimize the reconstruction loss only it doesn't really care how we compress the vector as long as it's able to reconstruct the input as good as possible and so it learns some arbitrary function to compress and decompress data this is fine for training but it's a problem when generating music during generation we don't have the encoder and we're working with just a decoder since we don't know the function our auto encoder learned chances are if we pass in a compressed vector to the decoder this is going to generate gibberish particularly I'm talking about the last leg of the flow and that's why the anonymous rapping sounded so Shh or beautiful this is where variational auto-encoders are better like normal auto-encoders they minimize a reconstruction loss but with the added constraint that it learns to do so with a specific distribution typically like a standard Gaussian distribution in other words we know how the compression and decompression is happening so during the music generation phase we don't really have the encoder and we can sample a vector from the standard Gaussian distribution pass it to the decoder and it will give us some meaningful audio so replacing our auto encoder with a variational auto encoder let's see what we get with our jukebox hey jukebox have the anonymous celebrity rap to lose yourself by Eminem okay getting better are via learn to convert a continuous distribution of vectors to a sound but not all of these vectors produce meaningful and clear audio and this is one of the reasons why some muddled nests and unclarity is still there in our generated audio to remedy this we can have our VA e to learn to encode a specific set of discrete vectors to a clear sound this is the basis for vector quantized variational auto-encoders during training we pass in a raw audio the encoder compresses it to a vector and we determine the closest discrete vector and then decode that as the sound and during generation the output of our transformer will be mapped to one of these discrete vectors and we can get a meaningful audio output so the vector quantized V II's can get clearer sounds but we can make further improvements on this by using multiple vqv AES at different compression levels this draws inspiration from a study on hierarchical vector quantizer AES used in image compression the goal is to generate even more realistic images by using a hierarchy of two vqv AES so we have a top v QV AE and a bottom v QV AE both of these learn different compressed representations of the image the top-level vqv a II only learns about the global information of the image like the contours and the large strokes but the bottom level vqv AE is larger and it learns about local pixel information like texture shading and color gradients and both of these learn representations are fed into the bottom level decoder to reconstruct the original image the bottom representation is conditioned on the top representation so that it doesn't need to learn the entire representation itself it makes the best use of its space leading to stunning images during the generation phase so like this we have three levels of VQ v's that we train in our case atop a middle and a bottom the top is the most compressed and the bottom is the least compressed but we don't feed the top to the middle in the middle to the bottom like we saw in the image case this is because the researchers observe that the top and the middle were passing all the information to the bottom and that made them useless it's it was similar to just having a single vq via e and so they just trained three v QV II's separately and in parallel it's interesting how we went from auto-encoders to using variational autoencoders to vector quantizer variational autoencoders and then to a hierarchical structure of the same each of them making improvements to generate better and better audio here's the initial structure that we've built in past one but now this autom encoder is actually three v QV a ease with different compression levels but to have these representations interact with each other we introduced three transformers the top transformer takes in lyrics genre and artist information to give the top level compressed representation and the second transformer converts this to the middle compressed representation the third transformer converts it to the bottom compression representation and then we pass this into the bottom vqv a decoder that we trained previously to get a newly generated audio the song hey jukebox have the UH nonnamous celebrity rap to lose yourself by Eminem his palms are sweaty knees weak arms a heavy there's vomit on his sweater already mom's spaghetti he's nervous but on the surface he looks calm and ready to drop bombs but he keeps on forgetting this is the final architecture and the end of past - things are looking a lot more concrete now but how exactly are we training this thing so in past three let's start with the vq ve training and then move on to transformer training we start by training our 3 v QV II's in parallel raw audio is a continuous stream so we need to break it down into fixed sized chunks let's say that we're dealing with a 20 second piece of audio here the top layer we'll break it down into five chunks the mid layer we'll break it down into 10 chunks and the bottom layer let's say it breaks it onto 20 chunks we pass the chunks one at a time through the encoder to get these individual colored bars note that each color bar has the same vector dimensions despite the width being different this top blue bar for example represents like four seconds of audio encoded into a 64 dimensional vector this first purple bar in the middle is the first 2 seconds of audio encoded into a 64 dimensional vector this first bottom brown bar is just one second of audio that's encoded into a 64 dimensional vector passing in all these chunks through the encoder gives us the compressed representations now we perform vector quantization for each of these colored pellets we determine the closest codebook vector the codebook is a list of vectors the blue one here ran in the top level is closest to vector 3 the purple one is closest to the 5 and this magenta one is closest to the 4th and for the codebook lookup we replace each of these numbers with the actual corresponding codebook vector and then we just decode each vector one at a time to get back the audio chunks and stitch that to get the original signal during this training we want to minimize the reconstruction loss we want to learn these codebook vectors and we also have a commit loss to stabilize the encoder so that's the bulk of this hierarchical vqv a training now once all these are trained we train our transformers we take a piece of audio pass it through our three vqv AES to get the top middle and bottom vector quantized representations these are then used to train our transformers the Transformers have an encoder decoder architecture the top level transformer takes in the lyrics artist genre and other conditional information to generate some intermediate vectors the decoder then takes these vectors and a start token to just generate the first highly compressed vector that represents a part of your generated song this is compared to the top level vq v AE vectors to minimize the difference during training time and it does this sequentially generating one vector at a time by our decoder transformer now that's a top level transformer moving on to the mid level transformer this takes one of the top level vectors of the vqv a II take some lyrics some genre and artist information all to generate some encoding vectors and the decoder will use this to generate two vectors generated one at a time during training these vectors are compared to the mid-level vqv a representation to minimize the difference the input vector should represent two seconds of data but the transformer converted it into two vectors that represent one second of data each I say - assuming that the compression rate of the mid level v QV a is two times that of the top level in the actual implementation though the says four times as much but I think you get the idea the bottom level transformer works in a similar way it takes a vector generated from the mid-level vq v AE as input this along with lyrics and genre is encoded into vectors these are passed into the decoder to sequentially generate two vectors one at a time they are compared to the bottom vqv vectors to minimize that difference if the input corresponds to one second of audio the output would be two vectors corresponding to half a second of audio each and in this way we can pass in a number of audio clips - or transformers to train them sweet so now during generation time we would get lyrics genre and artist information all encoded into vectors pass it to our top level transformers to get a very compressed representations pass each of these to the mid level transformer to decompress and pass each of these to the bottom level transformer to further decompress and then pass each of these vectors to the bottom level vqv a decoder to generate audio chunks stitch them together and we get view music and that's it I hope this covered different levels of understanding there's still a lot of detail that I did leave out but if you made it this far I'll add some references down in the description below of research papers reference papers and blog posts that you can check out thanks for watching and I will see you soon bye bye