VQ-VAEs: Neural Discrete Representation Learning | Paper + PyTorch Code Explained
❤️ Become The AI Epiphany Patreon ❤️ ► https://www.patreon.com/theaiepiphany

In this video I cover VQ-VAEs papers:
1) Neural Discrete Representation Learning
2) Generating Diverse High-Fidelity Images with VQ-VAE-2 (the only difference is the existence of a hierarchical structure of latents and priors)

Many novel interesting AI papers such as DALL-E and Jukebox from OpenAI as well as VQ-GAN build off of VQ-VAEs, so it's fairly important to have a good grasp of how they work.

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 
✅ VQ-VAE1 paper: https://arxiv.org/abs/1711.00937
✅ VQ-VAE2 paper: https://arxiv.org/abs/1906.00446
✅ PyTorch code: https://colab.research.google.com/github/zalandoresearch/pytorch-vq-vae/blob/master/vq-vae.ipynb
✅ ELBO explained: https://mbernste.github.io/posts/elbo/

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 

⌚️ Timetable:
00:00 Intro
01:10 A tangent on autoencoders and VAEs
07:50 Motivation behind discrete representations
08:25 High-level explanation of VQ-VAE framework
11:20 Diving deeper
13:05 VQ-VAE loss
16:20 PyTorch implementation
23:30 KL term missing
25:50 Prior autoregressive models
28:50 Results
32:20 VQ-VAE two

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 
💰 BECOME A PATREON OF THE AI EPIPHANY ❤️ 

If these videos, GitHub projects, and blogs help you,
consider helping me out by supporting me on Patreon! 

The AI Epiphany ► https://www.patreon.com/theaiepiphany
One-time donation: 
https://www.paypal.com/paypalme/theaiepiphany 

Much love! ❤️


Huge thank you to these AI Epiphany patreons:
Eli Mahler
Petar Veličković
Zvonimir Sabljic

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 

💡 The AI Epiphany is a channel dedicated to simplifying the field of AI using creative visualizations and in general, a stronger focus on geometrical and visual intuition, rather than the algebraic and numerical "intuition". 

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 
👋 CONNECT WITH ME ON SOCIAL 
LinkedIn ► https://www.linkedin.com/in/aleksagordic/ 
Twitter ► https://twitter.com/gordic_aleksa 

Instagram ► https://www.instagram.com/aiepiphany/ 
Facebook ► https://www.facebook.com/aiepiphany/ 

👨‍👩‍👧‍👦 JOIN OUR DISCORD COMMUNITY:
Discord ► https://discord.gg/peBrCpheKE

📢 SUBSCRIBE TO MY MONTHLY AI NEWSLETTER:
Substack ► https://aiepiphany.substack.com/

💻 FOLLOW ME ON GITHUB FOR COOL PROJECTS: 
GitHub ► https://github.com/gordicaleksa 

📚 FOLLOW ME ON MEDIUM:
Medium ► https://gordicaleksa.medium.com/
▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬

#vqvae #discretelatents #generativemodeling
 What's up? In this video I'm covering this older paper called Neural Discrete Representation Learning by Aaron van den Oord, Oriol Vinales and Corey... Yes. From DeepMind. Anyways, they introduced this model called VQVAE and the reason I'm covering it is because it's a seminal paper and many of the novel research coming from the AI field, like such as Dali and Jukebox and VQGAN, all of them are using VQVAE as a building block and so I think it's really important to understand this paper. So aside from diving deep into the paper as I always do, I'll also show you and walk you through a code snippet in PyTorch so you can really relate the code with the semantics of this paper. And if you found it useful, just share this video out and that's the best thing you can do to help me. Awesome, let's dig into it. So our model, the Vector Quantized Variational Autoencoder, differs from VAEs in two key ways. So the encoder network outputs discrete rather than continuous codes and the prior is learned rather than static. So let me just make a small digression into what autoencoders are and what variational autoencoders are before we can dig into VQVAEs, okay? So autoencoders are pretty simple. So you basically have something called an encoder network. So I'll just depict it like this. So it's E. You have something called like a latent code and latent is just a fancy way of saying hidden. So this is a latent vector here and then you have a decoder, okay? And let me maybe take an example of autoencoding images. So you take an image as the input, you project it down into this lower dimensional space. So maybe this is your ImageNet image, 224 times 224 times 3. And you may be projected down to, I don't know, maybe 30 dimensions here. And then you have to reproject it back into the original image space. So this and make sure that the reconstruction error is as low as possible. So you just basically make sure that this image is the same as this one, even though you kind of probably lost some information. Encoding this into this lower dimensional space. And you usually do that like using MSC loss. These are your continuous value images like from 0 to 1, you'll be using MSC. Okay, so that's autoencoder. And basically what variational autoencoders do additionally is they instead of having this deterministic bottleneck layer, they make it stochastic. So what encoder in VAE does is it will parameterize a distribution instead of a static latent vector. So usually people use isotropic Gaussians. That doesn't have to be that way. But like that's because Gaussians have really nice computational properties. And so people usually use those. But don't think that VAE have to use Gaussians. They are much more abstract framework than just using Gaussians. So what you'll have usually is you'll have encoder and you'll be outputting, for example, 30 dimensional vector of means. And you'll be outputting 30 dimensional vector of standard deviations. And you usually do log of standard deviation because you want to make sure that standard deviation is positive. And because it's like an argument of a log function, that means domain is greater than zero. So yeah, everything is computationally nice here. So once you have U's and sigma's, what you do instead of just passing it through the decoder, you sample a vector from this distribution and then you feed it back into the decoder. And so now the difference, so when you train these VAEs, you have two terms basically. And basically you also have the reconstruction term. Let me call it R. So that's that part. And you also have something called like you're doing a KL divergence between the prior and between whatever distribution you've outputted here. And let me kind of try and break that down a bit. So you can treat the KL term as a regularization term. But there is a deeper probabilistic explanation for why that is and why are we using KL instead of something else. So basically the reason is we are trying to maximize this thing called ELBO or evidence lower bound. And it turns out that like maximizing ELBO is the same thing as basically trying to minimize reconstruction loss, plus trying to minimize the KL divergence between. So this thing here is basically your approximate posterior. So this is Q of Z. So you're trying to infer the Z, the latent code, given your input, right? So given the X. So that may be your image. And what that second term in VAE is, you're trying to just minimize the KL between this. So this is your approximate posterior and the prior. And you usually put prior to be like, as I said, Gaussian. And so you can treat this either as a regularization. So that would be like your neural network kind of a perspective. Or you can treat this as a consequence of variational inference, where basically the reason you're doing this is you're trying to make your posterior as close as possible to the true posterior, which would be P of Z given X. And after like doing some math, you end up maximizing ELBO, which boils down at the end to just maximizing, to minimizing reconstruction loss and minimizing this KL divergence between your prior, which is usually Gaussian, because as I said, it's computationally convenient, and this approximate posterior. So all in all, the reason you're doing this is because you're trying to use these as generative models. So that means you want to be in a situation where you can just plug in some latent vector here and out comes some novel image. And the reason you can do this with these ordinary autoencoders is because their latent space doesn't have any constraints. So you basically, if we were to project this 30 dimensional vector into a 2D space, we'd end up with something like this. So maybe your cat image would be like projected here, your dog image would be projected here, some other dog would be projected here, and then you'd have like some like car, whatever. And the thing is, if you tried and queried, if you took a random latent vector like here and tried and feed it through the decoder, you'd probably get a totally random function, like a totally random image. And what you want to do instead, you want to impose certain structure into this latent space. So you want to make sure that it's continuous. So you want to make sure that if you're interpolating between these two points, you're getting the meaningful images all the way down from this point A all the way down to this point B. And VAEs do exactly this because of this KL term here, they are regularizing the latent space and making it much more continuous and meaningful compared to your standard autoencoders. So that was your 101 for autoencoders and VAEs. Now let's see basically how these VQ VAEs work. Getting back here, so they mentioned VAEs have continuous space. So as you can see here, these latent vectors will be continuous, whereas we'll soon see that the VQ VAE has a discrete latent variables. On the other hand, so the prior is static here. So as I said, so your prior here is a standard like Gaussian. And here like in VQVs, as we'll soon see, we'll be training some autoregressive model to in order to learn the prior. And those are some of the differences and we'll see those in a minute. OK, so we concentrate on discrete representations, which are potentially a more natural fit for many of the modalities we are interested in. Language is inherently discrete. Similarly, speech is typically represented as a sequence of symbols and images can often be described concisely by language. Furthermore, discrete representations are a natural fit for complex reasoning, planning and predictive learning. So that's some motivation for why we want to discretize the latent space. They also avoid something called posterior collapse by doing this. And that's additional benefit of using quantized VAEs. OK, let me let me jump into this diagram and explain the method and then we'll slowly start kind of this like breaking down the components into smaller pieces. So on a high level, again, we have this autoencoder like high level structure, as you can see. So we have the image, we encode it using some CNN into into this representation here. And then what we do is so this is the part that's peculiar to VQVs. So what we do is we take these vectors. So we take this one, for example, and what we do is we snap it onto one of these vectors from this thing called a codebook. So this thing here is called a codebook and it contains vectors which are also learnable and trainable. And so we want to find the closest one where closeness is defined by the L2 norm. And once you do that, you can basically map this one. So maybe the closest one was maybe this one, vector one. And so you'd put one here, as you can see in this symbolic table. And then what you do is you basically index using these indices into the codebook and you extract these like discrete latent vectors. And then only then do you feed those into the decoder and out comes the image, which is hopefully perfectly reconstructed version of your input image. So that's the high level overview. And as you can see, one difficulty with this is how do you actually back prop? So the forward prop is pretty straightforward, but the back prop is kind of difficult because this process of snapping onto the nearest like codebook vector is not differentiable. So what they instead do is something called straight through gradient. Like basically, as you can see, this shape here has the same shape as this thing here. So what I do is whatever the gradients for this volume here are, they just copy paste them into the encoder and then you can back prop through the encoder as usual. So the only part that's tricky is to just basically copy paste gradients here over to here from the decoder to the encoder. And we'll see like don't worry. We'll see the details in the code later as well. So the intuition behind this is that if if if you're encoding vectors, so if this vector here is encoded really closely to this codebook vector. So as you can see, so if this was a gradient for the codebook vector, we can expect that by moving the encoded vector a bit to the right. So the next time we feed forward that image will maybe snapped to this vector and will lower the reconstruction loss by doing that. So that's some rough intuition why this works. So hopefully the assumption is that is that these encoded vectors are close to the codebook vectors. And if that's the case, then these gradients kind of also are good gradients for those encoding vectors. And then you can just lower the the the basically the reconstruction loss. OK, so that's the high level picture. Now, let me kind of slowly start kind of breaking down these equations. And it's pretty easy. So this is just formalized version of what I've just explained verbally. So this is your encoded vector. This is your codebook vector. And you're trying to find that codebook vector, which is the closest to this encoded vector where closeness is defined using L2 norm, as you can see by these brackets and number two here. So once we find that index, that's our like that's the so we have this is our approximate posterior. And as you can see, it's basically one hot is deterministic distribution of our input variable. OK, so this just says the same thing I just explained. So once you find that K, you basically assign to so you assign to this Zeddy, you actually assign it as a proxy. You put this Zedq, which is just the closest vector to that Zeddy. And that's what we do in this discretization step. So we find the closest one and we basically use so we'll be using this E1 vector instead of this vector here. That's the idea of this equation here. So note that there is no real gradient defined for equation two. However, we approximate the gradient similar to the straight through estimator and just copy gradients from decoder input here to encoder output Zeddy. And that's the this this this line, this red line, they draw here. OK. And since the output representation of the encoder and the input, so this is the output representation of the encoder. So and the input of the decoder. So this is this one share the same dimensional space. The gradients contain useful information for how the encoder has to change its output to lower the reconstruction loss. So those are some of the things I already explained. OK. So finally, this is how the loss looks like. So we contrast this to VAE where we have KL and we have reconstruction loss. So here we have also we have the reconstruction loss. So that's this log P of X given Zed. So this is your likelihood. And you have these two terms. So first, let me kind of explain this one and explain why do we get MSC when we have these like weird log likelihood equations here. So I guess this can confuse some like like people who are who are new to this channel. So this is a likelihood. And the reason we get MSC out of this is the following. So we assume that the likelihood is Gaussian and under that assumption, which is reasonable if you have if you're trying to reconstruct images where the pixels are continuous. Under that assumption, if you just replace this P with your like Gaussian, like you have some constant, right? You have some e raised to the power of blah, blah, blah. You have some square here. And as you can see, because of the log, we'll just be ignoring constant E and log will be canceled out because you're inverse of each other. And you'll just be and you'll just end up with this square term. And that's why you have square loss as the as the optimal loss term under the assumption that your data, the likelihood of your data is Gaussian. So, yeah, hopefully that was helpful. And the second two terms are interesting because so compared to V is we don't have KL here. We have these two. So as you just mean, stop gradient. So that means in Pytorch, you just be detaching that part from the computational graph to just freeze those and you try and push these vectors. So these are the codebook vectors. You push them to their corresponding encode encoded vectors. And the same thing here. You basically freeze the codebook vectors and you're pushing the encoded vectors towards those codebook vectors. So that's that's how the loss looks like. It's fairly simple. We'll see that in the code in a couple of minutes. And so what I say here is so decoder optimizes the first loss term only the encoder optimizes the first and the last loss term loss terms. And the embeddings are optimized by the middle last term. So let's see why that is. So first things first, you can see we learned the embeddings by this second last term. So that's what I said. We'll learn encoder through both this term. As you can see, we are pushing the output from the encoder here towards these frozen codebook vectors. And we're also using reconstruction loss because, as you remember, we have that straight through gradient. So we are also training encoder using this reconstruction term, whereas the decoder is only trained. So those weights are only trained using this reconstruction loss. And these two terms do not contribute to its updates. OK, that was pretty much detailed explanation of how how it works. Now, I thought kind of walking you through an example of how this looks in PyTorch. And hopefully it will be easier to understand. And once you connect these images and like paper with concrete code. So I usually only cover code, only cover papers in my video. So now I thought doing this as an experiment. And like, let me know what you think down in the comments. OK, so so there is this and I'll link the code in the description, obviously. So there is this vector quantizer class. And let's see how it works. So we have embedding dimension and we have number of embeddings. So number of embeddings is this thing here. So it's basically this is your number of embeddings, the number of vectors in your code table and the embedding dimension. And the embedding dimension is this dimension here. OK, so once we have that, we just create this embedding object in PyTorch and we just initialize it uniformly. Blah, blah, that's that's easy. And we have this commitment cost part, which is basically this beta term. So that's just a weight with this commitment term of the loss. OK. And now let's see how the forward function looks like. So first, let me just give you some context. So let's let's take an as an example from this snippet. So as an example for a BCHW tensor of shape 16, 64, 32, 32, we first converted to blah, blah, blah. So what happens here is you take 16 images. So you feed 16 images into your decoder. Out comes this tensor here. OK, so we still have 16 images. But what changed is that the number of channel as you probably know, like CNN's usually you start with three channels and then you're as you're going through the CNN layers, you get more and more channels. But the spatial dimension kind of diminishes. OK, so that's what we see here. So we went from three to sixty four and the height and width went to thirty two by thirty two. OK, so so what I first do is they converted into this BCHWC vector instead. And then they flatten it into this representation here. And all of these 16000 vectors of size 64 will be quantized independently. So they'll be snapped to those codebook vectors. OK. And in other words, the channels are used as the space in which to quantize. All other dimensions will be flattened and be seen as different examples to quantize. 16,384 in this case. So that's in context. And now let's see the code. So this is the first. So this is the permutation that happens. So we want to get it into this format contiguous. Just make sure that in the memory, all of these memory locations are contiguous. Just some low level low level memory management stuff. So here is the flatten part. So we converted into the representation of 16000 by 64. And so here is this distances calculation. So this is a part that we use to find what what are the closest codebook vectors. And you can go through the like every single symbol here, but I'm just going to like do it high level. Basically. So flat input is it has shape. Let's say n times 64. OK, so n is 16000 in this example. But let me just use the symbol and instead. So and we have embedding table, which is as you remember, K and 64. So it's a couple K, 64. OK, so what we'll end up having here is basically we'll get a table that's n, K. And that means for every single encoded vector, we'll find its distance to all of the codebook vectors. And that's why we have K here. And once we have this, so once we have this distances array, what we do is we do simple argument over dimension one. So that means let me just kind of pictorially draw this. So we have, as I said, distances like is a matrix like this. And what we do is we focus on a single row and we find the lowest distance here. So that's the argument over dimension one. And we find, let's say that's maybe number three. So maybe this vector here is the closest one and that's maybe number three. So we'll map this into a novel vector, which will have number three here. So that's your symbol. So that's this part here. Let me just go back here. So those are whoops, my God. OK, so that will be your. So this vector here mapped into one. OK, here we just had another example with number three, but that's it. So once we do that, what they do additionally is once you create this vector of indices, you just cast it into a bit different representation. So instead of having three, you'll do it like this. You'll you'll basically expand it into a matrix and times K again. And this time you're going to create one hot vector. So you're going to put zero, one, two, three. So you want to put one here and all of the zeros will be here. So this is just a different representation of this very same information. OK, so that's the encodings data structure. And why do we do that is because we can then multiply this encodings data structure with the embedding table. So, again, embedding table is the code codebook table. So you'll have just a second. So we have K here. So this will be K dimensional and this will be 64. So these are this is your codebook vector code table. So if you can see here, just by multiplying this encodings data structure with the codebook vector, you'll end up extracting vector number three from from here. And that's basically the quantization process we're talking about. So you'll end up having basically quantized vectors. You'll have you'll end up having if you multiply these, you can see you end up with n times 64 dimensional table. And those will contain the quantized vectors. So that's how the quantization is done. Simple matrix multiplication. And finally, the interesting stuff here is so this E latent loss is this term here. So that's this term here. As you can see, we put a stop gradient onto the codebook vectors. And as you can see here, so these are the codebook vectors. We put a stop gradient, which is implemented as detached in PyTorch. We do the same thing for this second term. So this is this part where we put stop gradient on the encoded vectors. So encoded vectors are called inputs in our code. And we just do detach. And that's it. Just simple MSC laws there. And again, this is beta. So that's your loss. OK. And now for the funny part, and I'll just ignore this, is how do we actually implement these straight through gradient? So this is the way we implement that estimator. So if you take a look at this equation, obviously inputs and inputs will kind of cancel out and you end up with quantized. So the reason they added this part here is because when we do backprop, this part is basically excluded from the computational graph. And so we end up with inputs. And so that means that basically whatever the gradients you get for your outputs, for your quantized variable here, those will just be copy pasted into the inputs. And that's the part I just explained pictorially here. Basically, you have whatever the gradients here are. You're just going to copy paste them into here. And once you have that, then you can just do because everything is differentiable in the encoder. You can figure out the gradients for every single weight and then you can just update your network. OK. So you may wonder why we are not using the KL term. Let me try and give you some intuition there. So they say here, since we assume a uniform prior for Z, so compare that to VAEs, which had like isotropic Gaussians as the basically standard Gaussians as a prior. Here we have we assume uniform prior for Z. The KL term that usually appears in the elbow. So this is your evidence lower bound is constant with respect to the encoder parameters and thus can be ignored for training. So let me try and digest this for you a little bit. So that's so pure. If that if we take a look at that, we have a uniform distribution here. OK. So so this will have K. Points in the domain here and what we'll end up with. So how our posterior looks like, so how our approximate posterior Z given X looks like is it's as I mentioned, it's a deterministic function pretty much. You will have you'll end up having a single vector that's the closest to your encoded vector. So you'll end up with something like this. And if you take a look at what you get by doing a KL divergence between these two. So if you do KL between these two, you'll end up with basically just one times log. Then we have in the numerator, we'll have one and in the denominator we'll have we'll have because this is uniform, we have K points here. This is going to be one over K probability. So the probability that corresponds to this one, let's maybe that's this one, will be one over K. So you end up with KL divergence that's constant and that's equal to log of K. OK. And so you basically can't minimize it because it's constant. And they mentioned it somewhere here. So our proposal distribution blah, blah is deterministic. And by defining a simple uniform prior of Z, we obtain a KL divergence constant and equal to log of K. So that's why they can't minimize it because it's constant in this setup. And yeah, I just wanted to explain that a little bit more. OK. So for the prior part, they say here, whilst training the VQVAE, the prior is kept constant and uniform, as we just saw. After training, we fit an autoregressive distribution over Z, P of Z, so that we can generate X via ancestral sampling. So all of this sounds maybe a bit scary. Let me try and demystify this part. So if you're familiar with language modeling, if you watch my GPT-3 video, I'll link it somewhere here or like transformers, whatever, if you're familiar with language modeling in general, this will be really easy. So what you basically do here is in a way you just try to predict the next token. So in the case of language modeling, you're trying to predict a token that represents some subword or something or like a word or a character. Here, you're just trying to predict the next symbol. But like the semantics of the process is the same. So imagine you flatten this into a sequence and you have the same, like basically same setup as for language modeling. So let's say we have a simple example, not like a not like a matrix of symbols, but maybe just like a three symbols. Like let's say we have three, two, five. OK. And how you train this all regressive model is the following. So you'd input something called some special symbol called start of a sequence. Then you'd input three here and then you'd input two here. And as you can see, you're trying to predict. So given the special token, you're trying to predict three. Then given three, you're trying to predict two because that's the next token, as you can see here. And now given two, you're trying to predict five because that's the next token in the sequence, as you can see. So this goes as an input to your model. These are the targets. And basically you're just you'll you'll basically use some like cross entropy and you'll backprop through the weights of the model. And that's how you're going to train this this this token predictor model. And it will going to be all regressive, which means once you output a token here, you're going to feed it back to the input and then you're going to keep on generating like that. So once you have it trained on this discrete space that was so basically there are two things you do here. First things first, you basically train this VQVAE end to end. So to learn how to reconstruct images. And then you have you basically have a codebook that's that's fully trained and you have encoder and decoder, which are fully trained. You just freeze all of those and you just train this predictor model. And then when you want to generate new novel images, what you do is you just prompt the the all regressive model with the special token and it will just output some number. Let's maybe say three and you feedback the three inside of the input of your all regressive model and it will generate something else like maybe four. And you do that all the way until you have enough elements in your in your like latent space. So that will be thirty two by thirty two in this example here. And once you have that, you just feed that through this frozen decoder and out comes some novel image. So that's the explanation of how this whole system works together. Hopefully you you now understood it a bit better. OK, so having demystified that, hopefully let's see some results. And first thing they mentioned here is that they can compress their data by a lot. And that's hopefully obvious. But like so they say here in this experiment, we show that we can model X, which is hundred twenty eight by hundred twenty eight in their example by compressing them to thirty two by thirty two discrete space with K equals five hundred twelve. So case the size of that codebook codebook table, why are purely the convolutional PX given set. So decoder. So a reduction of and you can see how much here and basically nine is here because if you have a naive encoding method of your symbols, you'll just have nine bits. You need nine bits in order to represent five hundred twelve symbols if you assume that all of the code like sizes are fixed. OK, and basically you can see here some reconstructions. And what I want you to notice here is because this will be important for the later Dali paper. I'll create like a paper over. I'll do and VQ again as well. You can see that there is some blurriness to these images. If you take this dog here and you'll see like the reconstruction is pretty good, but there is some blurriness happening here. It's especially noticeable here. And for this real image here, you can see it's much blurrier. You can take a look at the paper yourself. But that's that's something I want you to notice here. So those were the reconstructions. And here after we do that sampling process in the inside of the latent space using that order aggressive model, you can sample novel unconditional images and you can see some results here. And you can see it's really diverse. And that's one of the advantages these likelihood models have compared to GANs. They cover all of the modes because of the way they are trained, whereas GANs sometimes can omit some modes or they can even like go into something called mode collapse where they basically output just a single image or similar looking imagery. OK, they also do this not only for images, they can also compress and reconstruct audio, video, et cetera. So I found the audio part really interesting. So this is your original audio after you reconstructed. You can see the waveforms are not quite the same. But what's interesting is that they have the same content. Couple of nice remarks they made here is it can be heard from the provided samples and as shown in Figure 7 reconstruction. So Figure 7 is this one here. You can find the samples on the website. You can hear them at your own pace. But like the reconstruction has the same content, same text contents, but the waveform is quite different and prosody. So the way you accent your sentence in the voice is altered. So this means that the VQVAE has without any form of linguistic supervision, so without any supervision, learned a high level abstract space that is invariant to low level features and only encodes the content of the speech, which is awesome since we don't have any labels here. So while samples drawn from even the best speech models like the original WaveNet sound like babbling, samples from VQVAE contain clear words and part sentences. So what I recommend you do is you go ahead and listen to the samples from this model and from WaveNet, which was very popular, like parallel WaveNet, like that was a full up work to WaveNet, was actually used in production at Google. It's maybe still used. I'm not sure, but I know it was used in production for Google Assistant. OK, so I think that's pretty much it for this paper. Let me just briefly go and explain you how this VQVAE 2 version differs from this version 1. So first thing you can notice is it's got much better reconstructions. You can generate much better imagery using this model. The main difference is so everything remains the same. The loss remains the same. The only difference is they have this hierarchical structure of the of the latents as well as the priors. So what I do instead is they have this top level codebook vectors and they have this mid level codebook vectors and they are of different resolutions. So that means this one will have to focus on some like higher level structure of the image. This one can maybe even focus on some details, texture, et cetera. And once you combine them into decoder, you can reconstruct much higher quality imagery. Same thing here. You train this autoregressive model in this latent space and then you use that to condition the next level prior autoregressive model to basically regress this new latents and then you can decode them into imagery. So that's the only difference pretty much between this model and the last one. So they're using this hierarchical structures. And as you can see, they get much better imagery. Like this is even compared like comparing even this to, I don't know, style again. V2, the images look really, really cool, although style again can can generate much higher resolution imagery. I think so. But yeah, either way, this is really cool for a VAE. Usually GANs have much better reconstruct image images than VAEs. OK, here are some more images. Again, I think this is class conditional on ImageNet. I think doesn't matter that much. Let me just see. Yeah, these are some high resolution images. You can see they're pretty good. You can see it's still not perfect. There is some distortions here happening. I don't know. Maybe I'm wrong. Hopefully you found this video insightful. Let me know what you think down in the comments, whether you like this format, whether the code part was useful. If it was, I'm going to try and include more of these into my next videos. So until next time, subscribe and share this video and bye bye.
[{"start": 0.0, "end": 8.52, "text": " What's up? In this video I'm covering this older paper called Neural Discrete Representation Learning by Aaron van den Oord, Oriol Vinales and Corey..."}, {"start": 10.32, "end": 10.96, "text": " Yes."}, {"start": 11.36, "end": 12.88, "text": " From DeepMind."}, {"start": 13.040000000000001, "end": 21.080000000000002, "text": " Anyways, they introduced this model called VQVAE and the reason I'm covering it is because it's a seminal paper"}, {"start": 21.080000000000002, "end": 28.28, "text": " and many of the novel research coming from the AI field, like such as Dali and Jukebox and VQGAN,"}, {"start": 28.28, "end": 34.68, "text": " all of them are using VQVAE as a building block and so I think it's really important to understand this paper."}, {"start": 34.68, "end": 42.32, "text": " So aside from diving deep into the paper as I always do, I'll also show you and walk you through a code snippet in PyTorch"}, {"start": 42.32, "end": 46.96, "text": " so you can really relate the code with the semantics of this paper."}, {"start": 46.96, "end": 51.52, "text": " And if you found it useful, just share this video out and that's the best thing you can do to help me."}, {"start": 51.52, "end": 53.32, "text": " Awesome, let's dig into it."}, {"start": 53.32, "end": 59.8, "text": " So our model, the Vector Quantized Variational Autoencoder, differs from VAEs in two key ways."}, {"start": 59.8, "end": 66.76, "text": " So the encoder network outputs discrete rather than continuous codes and the prior is learned rather than static."}, {"start": 66.76, "end": 75.24000000000001, "text": " So let me just make a small digression into what autoencoders are and what variational autoencoders are before we can dig into VQVAEs, okay?"}, {"start": 75.24000000000001, "end": 81.24000000000001, "text": " So autoencoders are pretty simple. So you basically have something called an encoder network."}, {"start": 81.24, "end": 89.24, "text": " So I'll just depict it like this. So it's E. You have something called like a latent code and latent is just a fancy way of saying hidden."}, {"start": 89.24, "end": 93.24, "text": " So this is a latent vector here and then you have a decoder, okay?"}, {"start": 93.24, "end": 104.24, "text": " And let me maybe take an example of autoencoding images. So you take an image as the input, you project it down into this lower dimensional space."}, {"start": 104.24, "end": 112.24, "text": " So maybe this is your ImageNet image, 224 times 224 times 3."}, {"start": 112.24, "end": 116.24, "text": " And you may be projected down to, I don't know, maybe 30 dimensions here."}, {"start": 116.24, "end": 120.24, "text": " And then you have to reproject it back into the original image space."}, {"start": 120.24, "end": 125.24, "text": " So this and make sure that the reconstruction error is as low as possible."}, {"start": 125.24, "end": 133.24, "text": " So you just basically make sure that this image is the same as this one, even though you kind of probably lost some information."}, {"start": 133.24, "end": 137.24, "text": " Encoding this into this lower dimensional space."}, {"start": 137.24, "end": 146.24, "text": " And you usually do that like using MSC loss. These are your continuous value images like from 0 to 1, you'll be using MSC."}, {"start": 146.24, "end": 157.24, "text": " Okay, so that's autoencoder. And basically what variational autoencoders do additionally is they instead of having this deterministic bottleneck layer, they make it stochastic."}, {"start": 157.24, "end": 164.24, "text": " So what encoder in VAE does is it will parameterize a distribution instead of a static latent vector."}, {"start": 164.24, "end": 169.24, "text": " So usually people use isotropic Gaussians. That doesn't have to be that way."}, {"start": 169.24, "end": 173.24, "text": " But like that's because Gaussians have really nice computational properties."}, {"start": 173.24, "end": 178.24, "text": " And so people usually use those. But don't think that VAE have to use Gaussians."}, {"start": 178.24, "end": 181.24, "text": " They are much more abstract framework than just using Gaussians."}, {"start": 181.24, "end": 192.24, "text": " So what you'll have usually is you'll have encoder and you'll be outputting, for example, 30 dimensional vector of means."}, {"start": 192.24, "end": 196.24, "text": " And you'll be outputting 30 dimensional vector of standard deviations."}, {"start": 196.24, "end": 203.24, "text": " And you usually do log of standard deviation because you want to make sure that standard deviation is positive."}, {"start": 203.24, "end": 209.24, "text": " And because it's like an argument of a log function, that means domain is greater than zero."}, {"start": 209.24, "end": 213.24, "text": " So yeah, everything is computationally nice here."}, {"start": 213.24, "end": 218.24, "text": " So once you have U's and sigma's, what you do instead of just passing it through the decoder,"}, {"start": 218.24, "end": 225.24, "text": " you sample a vector from this distribution and then you feed it back into the decoder."}, {"start": 225.24, "end": 230.24, "text": " And so now the difference, so when you train these VAEs, you have two terms basically."}, {"start": 230.24, "end": 235.24, "text": " And basically you also have the reconstruction term. Let me call it R."}, {"start": 235.24, "end": 243.24, "text": " So that's that part. And you also have something called like you're doing a KL divergence between the prior"}, {"start": 243.24, "end": 249.24, "text": " and between whatever distribution you've outputted here. And let me kind of try and break that down a bit."}, {"start": 249.24, "end": 252.24, "text": " So you can treat the KL term as a regularization term."}, {"start": 252.24, "end": 259.24, "text": " But there is a deeper probabilistic explanation for why that is and why are we using KL instead of something else."}, {"start": 259.24, "end": 266.24, "text": " So basically the reason is we are trying to maximize this thing called ELBO or evidence lower bound."}, {"start": 266.24, "end": 272.24, "text": " And it turns out that like maximizing ELBO is the same thing as basically trying to minimize reconstruction loss,"}, {"start": 272.24, "end": 279.24, "text": " plus trying to minimize the KL divergence between. So this thing here is basically your approximate posterior."}, {"start": 279.24, "end": 285.24, "text": " So this is Q of Z. So you're trying to infer the Z, the latent code, given your input, right?"}, {"start": 285.24, "end": 294.24, "text": " So given the X. So that may be your image. And what that second term in VAE is, you're trying to just minimize the KL between this."}, {"start": 294.24, "end": 301.24, "text": " So this is your approximate posterior and the prior. And you usually put prior to be like, as I said, Gaussian."}, {"start": 301.24, "end": 309.24, "text": " And so you can treat this either as a regularization. So that would be like your neural network kind of a perspective."}, {"start": 309.24, "end": 319.24, "text": " Or you can treat this as a consequence of variational inference, where basically the reason you're doing this is you're trying to make your posterior"}, {"start": 319.24, "end": 324.24, "text": " as close as possible to the true posterior, which would be P of Z given X."}, {"start": 324.24, "end": 332.24, "text": " And after like doing some math, you end up maximizing ELBO, which boils down at the end to just maximizing, to minimizing reconstruction loss"}, {"start": 332.24, "end": 342.24, "text": " and minimizing this KL divergence between your prior, which is usually Gaussian, because as I said, it's computationally convenient, and this approximate posterior."}, {"start": 342.24, "end": 348.24, "text": " So all in all, the reason you're doing this is because you're trying to use these as generative models."}, {"start": 348.24, "end": 357.24, "text": " So that means you want to be in a situation where you can just plug in some latent vector here and out comes some novel image."}, {"start": 357.24, "end": 363.24, "text": " And the reason you can do this with these ordinary autoencoders is because their latent space doesn't have any constraints."}, {"start": 363.24, "end": 370.24, "text": " So you basically, if we were to project this 30 dimensional vector into a 2D space, we'd end up with something like this."}, {"start": 370.24, "end": 378.24, "text": " So maybe your cat image would be like projected here, your dog image would be projected here, some other dog would be projected here,"}, {"start": 378.24, "end": 385.24, "text": " and then you'd have like some like car, whatever. And the thing is, if you tried and queried, if you took a random latent vector like here"}, {"start": 385.24, "end": 394.24, "text": " and tried and feed it through the decoder, you'd probably get a totally random function, like a totally random image."}, {"start": 394.24, "end": 398.24, "text": " And what you want to do instead, you want to impose certain structure into this latent space."}, {"start": 398.24, "end": 404.24, "text": " So you want to make sure that it's continuous. So you want to make sure that if you're interpolating between these two points,"}, {"start": 404.24, "end": 411.24, "text": " you're getting the meaningful images all the way down from this point A all the way down to this point B."}, {"start": 411.24, "end": 420.24, "text": " And VAEs do exactly this because of this KL term here, they are regularizing the latent space"}, {"start": 420.24, "end": 425.24, "text": " and making it much more continuous and meaningful compared to your standard autoencoders."}, {"start": 425.24, "end": 436.24, "text": " So that was your 101 for autoencoders and VAEs. Now let's see basically how these VQ VAEs work."}, {"start": 436.24, "end": 444.24, "text": " Getting back here, so they mentioned VAEs have continuous space. So as you can see here, these latent vectors will be continuous,"}, {"start": 444.24, "end": 450.24, "text": " whereas we'll soon see that the VQ VAE has a discrete latent variables."}, {"start": 450.24, "end": 457.24, "text": " On the other hand, so the prior is static here. So as I said, so your prior here is a standard like Gaussian."}, {"start": 457.24, "end": 466.24, "text": " And here like in VQVs, as we'll soon see, we'll be training some autoregressive model to in order to learn the prior."}, {"start": 466.24, "end": 469.24, "text": " And those are some of the differences and we'll see those in a minute."}, {"start": 469.24, "end": 476.24, "text": " OK, so we concentrate on discrete representations, which are potentially a more natural fit for many of the modalities we are interested in."}, {"start": 476.24, "end": 484.24, "text": " Language is inherently discrete. Similarly, speech is typically represented as a sequence of symbols and images can often be described concisely by language."}, {"start": 484.24, "end": 489.24, "text": " Furthermore, discrete representations are a natural fit for complex reasoning, planning and predictive learning."}, {"start": 489.24, "end": 494.24, "text": " So that's some motivation for why we want to discretize the latent space."}, {"start": 494.24, "end": 498.24, "text": " They also avoid something called posterior collapse by doing this."}, {"start": 498.24, "end": 502.24, "text": " And that's additional benefit of using quantized VAEs."}, {"start": 502.24, "end": 512.24, "text": " OK, let me let me jump into this diagram and explain the method and then we'll slowly start kind of this like breaking down the components into smaller pieces."}, {"start": 512.24, "end": 517.24, "text": " So on a high level, again, we have this autoencoder like high level structure, as you can see."}, {"start": 517.24, "end": 523.24, "text": " So we have the image, we encode it using some CNN into into this representation here."}, {"start": 523.24, "end": 527.24, "text": " And then what we do is so this is the part that's peculiar to VQVs."}, {"start": 527.24, "end": 529.24, "text": " So what we do is we take these vectors."}, {"start": 529.24, "end": 538.24, "text": " So we take this one, for example, and what we do is we snap it onto one of these vectors from this thing called a codebook."}, {"start": 538.24, "end": 544.24, "text": " So this thing here is called a codebook and it contains vectors which are also learnable and trainable."}, {"start": 544.24, "end": 550.24, "text": " And so we want to find the closest one where closeness is defined by the L2 norm."}, {"start": 550.24, "end": 553.24, "text": " And once you do that, you can basically map this one."}, {"start": 553.24, "end": 557.24, "text": " So maybe the closest one was maybe this one, vector one."}, {"start": 557.24, "end": 561.24, "text": " And so you'd put one here, as you can see in this symbolic table."}, {"start": 561.24, "end": 571.24, "text": " And then what you do is you basically index using these indices into the codebook and you extract these like discrete latent vectors."}, {"start": 571.24, "end": 580.24, "text": " And then only then do you feed those into the decoder and out comes the image, which is hopefully perfectly reconstructed version of your input image."}, {"start": 580.24, "end": 582.24, "text": " So that's the high level overview."}, {"start": 582.24, "end": 588.24, "text": " And as you can see, one difficulty with this is how do you actually back prop?"}, {"start": 588.24, "end": 599.24, "text": " So the forward prop is pretty straightforward, but the back prop is kind of difficult because this process of snapping onto the nearest like codebook vector is not differentiable."}, {"start": 599.24, "end": 602.24, "text": " So what they instead do is something called straight through gradient."}, {"start": 602.24, "end": 607.24, "text": " Like basically, as you can see, this shape here has the same shape as this thing here."}, {"start": 607.24, "end": 616.24, "text": " So what I do is whatever the gradients for this volume here are, they just copy paste them into the encoder and then you can back prop through the encoder as usual."}, {"start": 616.24, "end": 623.24, "text": " So the only part that's tricky is to just basically copy paste gradients here over to here from the decoder to the encoder."}, {"start": 623.24, "end": 627.24, "text": " And we'll see like don't worry. We'll see the details in the code later as well."}, {"start": 627.24, "end": 638.24, "text": " So the intuition behind this is that if if if you're encoding vectors, so if this vector here is encoded really closely to this codebook vector."}, {"start": 638.24, "end": 647.24, "text": " So as you can see, so if this was a gradient for the codebook vector, we can expect that by moving the encoded vector a bit to the right."}, {"start": 647.24, "end": 655.24, "text": " So the next time we feed forward that image will maybe snapped to this vector and will lower the reconstruction loss by doing that."}, {"start": 655.24, "end": 657.24, "text": " So that's some rough intuition why this works."}, {"start": 657.24, "end": 664.24, "text": " So hopefully the assumption is that is that these encoded vectors are close to the codebook vectors."}, {"start": 664.24, "end": 671.24, "text": " And if that's the case, then these gradients kind of also are good gradients for those encoding vectors."}, {"start": 671.24, "end": 675.24, "text": " And then you can just lower the the the basically the reconstruction loss."}, {"start": 675.24, "end": 683.24, "text": " OK, so that's the high level picture. Now, let me kind of slowly start kind of breaking down these equations."}, {"start": 683.24, "end": 689.24, "text": " And it's pretty easy. So this is just formalized version of what I've just explained verbally."}, {"start": 689.24, "end": 695.24, "text": " So this is your encoded vector. This is your codebook vector. And you're trying to find that codebook vector,"}, {"start": 695.24, "end": 703.24, "text": " which is the closest to this encoded vector where closeness is defined using L2 norm, as you can see by these brackets and number two here."}, {"start": 703.24, "end": 710.24, "text": " So once we find that index, that's our like that's the so we have this is our approximate posterior."}, {"start": 710.24, "end": 717.24, "text": " And as you can see, it's basically one hot is deterministic distribution of our input variable."}, {"start": 717.24, "end": 722.24, "text": " OK, so this just says the same thing I just explained."}, {"start": 722.24, "end": 731.24, "text": " So once you find that K, you basically assign to so you assign to this Zeddy, you actually assign it as a proxy."}, {"start": 731.24, "end": 739.24, "text": " You put this Zedq, which is just the closest vector to that Zeddy. And that's what we do in this discretization step."}, {"start": 739.24, "end": 746.24, "text": " So we find the closest one and we basically use so we'll be using this E1 vector instead of this vector here."}, {"start": 746.24, "end": 751.24, "text": " That's the idea of this equation here. So note that there is no real gradient defined for equation two."}, {"start": 751.24, "end": 760.24, "text": " However, we approximate the gradient similar to the straight through estimator and just copy gradients from decoder input here to encoder output Zeddy."}, {"start": 760.24, "end": 765.24, "text": " And that's the this this this line, this red line, they draw here. OK."}, {"start": 765.24, "end": 772.24, "text": " And since the output representation of the encoder and the input, so this is the output representation of the encoder."}, {"start": 772.24, "end": 778.24, "text": " So and the input of the decoder. So this is this one share the same dimensional space."}, {"start": 778.24, "end": 783.24, "text": " The gradients contain useful information for how the encoder has to change its output to lower the reconstruction loss."}, {"start": 783.24, "end": 786.24, "text": " So those are some of the things I already explained. OK."}, {"start": 786.24, "end": 794.24, "text": " So finally, this is how the loss looks like. So we contrast this to VAE where we have KL and we have reconstruction loss."}, {"start": 794.24, "end": 799.24, "text": " So here we have also we have the reconstruction loss. So that's this log P of X given Zed."}, {"start": 799.24, "end": 802.24, "text": " So this is your likelihood. And you have these two terms."}, {"start": 802.24, "end": 811.24, "text": " So first, let me kind of explain this one and explain why do we get MSC when we have these like weird log likelihood equations here."}, {"start": 811.24, "end": 817.24, "text": " So I guess this can confuse some like like people who are who are new to this channel."}, {"start": 817.24, "end": 822.24, "text": " So this is a likelihood. And the reason we get MSC out of this is the following."}, {"start": 822.24, "end": 832.24, "text": " So we assume that the likelihood is Gaussian and under that assumption, which is reasonable if you have if you're trying to reconstruct images where the pixels are continuous."}, {"start": 832.24, "end": 840.24, "text": " Under that assumption, if you just replace this P with your like Gaussian, like you have some constant, right?"}, {"start": 840.24, "end": 844.24, "text": " You have some e raised to the power of blah, blah, blah. You have some square here."}, {"start": 844.24, "end": 852.24, "text": " And as you can see, because of the log, we'll just be ignoring constant E and log will be canceled out because you're inverse of each other."}, {"start": 852.24, "end": 856.24, "text": " And you'll just be and you'll just end up with this square term."}, {"start": 856.24, "end": 866.24, "text": " And that's why you have square loss as the as the optimal loss term under the assumption that your data, the likelihood of your data is Gaussian."}, {"start": 866.24, "end": 869.24, "text": " So, yeah, hopefully that was helpful."}, {"start": 869.24, "end": 874.24, "text": " And the second two terms are interesting because so compared to V is we don't have KL here."}, {"start": 874.24, "end": 879.24, "text": " We have these two. So as you just mean, stop gradient."}, {"start": 879.24, "end": 889.24, "text": " So that means in Pytorch, you just be detaching that part from the computational graph to just freeze those and you try and push these vectors."}, {"start": 889.24, "end": 895.24, "text": " So these are the codebook vectors. You push them to their corresponding encode encoded vectors."}, {"start": 895.24, "end": 902.24, "text": " And the same thing here. You basically freeze the codebook vectors and you're pushing the encoded vectors towards those codebook vectors."}, {"start": 902.24, "end": 907.24, "text": " So that's that's how the loss looks like. It's fairly simple. We'll see that in the code in a couple of minutes."}, {"start": 907.24, "end": 918.24, "text": " And so what I say here is so decoder optimizes the first loss term only the encoder optimizes the first and the last loss term loss terms."}, {"start": 918.24, "end": 921.24, "text": " And the embeddings are optimized by the middle last term. So let's see why that is."}, {"start": 921.24, "end": 927.24, "text": " So first things first, you can see we learned the embeddings by this second last term."}, {"start": 927.24, "end": 931.24, "text": " So that's what I said. We'll learn encoder through both this term."}, {"start": 931.24, "end": 937.24, "text": " As you can see, we are pushing the output from the encoder here towards these frozen codebook vectors."}, {"start": 937.24, "end": 941.24, "text": " And we're also using reconstruction loss because, as you remember, we have that straight through gradient."}, {"start": 941.24, "end": 947.24, "text": " So we are also training encoder using this reconstruction term, whereas the decoder is only trained."}, {"start": 947.24, "end": 956.24, "text": " So those weights are only trained using this reconstruction loss. And these two terms do not contribute to its updates."}, {"start": 956.24, "end": 961.24, "text": " OK, that was pretty much detailed explanation of how how it works."}, {"start": 961.24, "end": 966.24, "text": " Now, I thought kind of walking you through an example of how this looks in PyTorch."}, {"start": 966.24, "end": 975.24, "text": " And hopefully it will be easier to understand. And once you connect these images and like paper with concrete code."}, {"start": 975.24, "end": 980.24, "text": " So I usually only cover code, only cover papers in my video. So now I thought doing this as an experiment."}, {"start": 980.24, "end": 988.24, "text": " And like, let me know what you think down in the comments. OK, so so there is this and I'll link the code in the description, obviously."}, {"start": 988.24, "end": 993.24, "text": " So there is this vector quantizer class. And let's see how it works."}, {"start": 993.24, "end": 999.24, "text": " So we have embedding dimension and we have number of embeddings. So number of embeddings is this thing here."}, {"start": 999.24, "end": 1004.24, "text": " So it's basically this is your number of embeddings, the number of vectors in your code table and the embedding dimension."}, {"start": 1004.24, "end": 1015.24, "text": " And the embedding dimension is this dimension here. OK, so once we have that, we just create this embedding object in PyTorch and we just initialize it uniformly."}, {"start": 1015.24, "end": 1022.24, "text": " Blah, blah, that's that's easy. And we have this commitment cost part, which is basically this beta term."}, {"start": 1022.24, "end": 1027.24, "text": " So that's just a weight with this commitment term of the loss. OK."}, {"start": 1027.24, "end": 1033.24, "text": " And now let's see how the forward function looks like. So first, let me just give you some context."}, {"start": 1033.24, "end": 1045.24, "text": " So let's let's take an as an example from this snippet. So as an example for a BCHW tensor of shape 16, 64, 32, 32, we first converted to blah, blah, blah."}, {"start": 1045.24, "end": 1053.24, "text": " So what happens here is you take 16 images. So you feed 16 images into your decoder. Out comes this tensor here."}, {"start": 1053.24, "end": 1067.24, "text": " OK, so we still have 16 images. But what changed is that the number of channel as you probably know, like CNN's usually you start with three channels and then you're as you're going through the CNN layers, you get more and more channels."}, {"start": 1067.24, "end": 1075.24, "text": " But the spatial dimension kind of diminishes. OK, so that's what we see here. So we went from three to sixty four and the height and width went to thirty two by thirty two."}, {"start": 1075.24, "end": 1083.24, "text": " OK, so so what I first do is they converted into this BCHWC vector instead."}, {"start": 1083.24, "end": 1092.24, "text": " And then they flatten it into this representation here. And all of these 16000 vectors of size 64 will be quantized independently."}, {"start": 1092.24, "end": 1098.24, "text": " So they'll be snapped to those codebook vectors. OK. And in other words, the channels are used as the space in which to quantize."}, {"start": 1098.24, "end": 1104.24, "text": " All other dimensions will be flattened and be seen as different examples to quantize. 16,384 in this case."}, {"start": 1104.24, "end": 1109.24, "text": " So that's in context. And now let's see the code. So this is the first. So this is the permutation that happens."}, {"start": 1109.24, "end": 1117.24, "text": " So we want to get it into this format contiguous. Just make sure that in the memory, all of these memory locations are contiguous."}, {"start": 1117.24, "end": 1123.24, "text": " Just some low level low level memory management stuff. So here is the flatten part."}, {"start": 1123.24, "end": 1130.24, "text": " So we converted into the representation of 16000 by 64. And so here is this distances calculation."}, {"start": 1130.24, "end": 1135.24, "text": " So this is a part that we use to find what what are the closest codebook vectors."}, {"start": 1135.24, "end": 1141.24, "text": " And you can go through the like every single symbol here, but I'm just going to like do it high level."}, {"start": 1141.24, "end": 1150.24, "text": " Basically. So flat input is it has shape. Let's say n times 64. OK, so n is 16000 in this example."}, {"start": 1150.24, "end": 1158.24, "text": " But let me just use the symbol and instead. So and we have embedding table, which is as you remember, K and 64."}, {"start": 1158.24, "end": 1169.24, "text": " So it's a couple K, 64. OK, so what we'll end up having here is basically we'll get a table that's n, K."}, {"start": 1169.24, "end": 1176.24, "text": " And that means for every single encoded vector, we'll find its distance to all of the codebook vectors."}, {"start": 1176.24, "end": 1184.24, "text": " And that's why we have K here. And once we have this, so once we have this distances array, what we do is we do simple argument over dimension one."}, {"start": 1184.24, "end": 1192.24, "text": " So that means let me just kind of pictorially draw this. So we have, as I said, distances like is a matrix like this."}, {"start": 1192.24, "end": 1197.24, "text": " And what we do is we focus on a single row and we find the lowest distance here."}, {"start": 1197.24, "end": 1202.24, "text": " So that's the argument over dimension one. And we find, let's say that's maybe number three."}, {"start": 1202.24, "end": 1206.24, "text": " So maybe this vector here is the closest one and that's maybe number three."}, {"start": 1206.24, "end": 1212.24, "text": " So we'll map this into a novel vector, which will have number three here."}, {"start": 1212.24, "end": 1216.24, "text": " So that's your symbol. So that's this part here. Let me just go back here."}, {"start": 1216.24, "end": 1221.24, "text": " So those are whoops, my God. OK, so that will be your."}, {"start": 1221.24, "end": 1229.24, "text": " So this vector here mapped into one. OK, here we just had another example with number three, but that's it."}, {"start": 1229.24, "end": 1238.24, "text": " So once we do that, what they do additionally is once you create this vector of indices, you just cast it into a bit different representation."}, {"start": 1238.24, "end": 1247.24, "text": " So instead of having three, you'll do it like this. You'll you'll basically expand it into a matrix and times K again."}, {"start": 1247.24, "end": 1254.24, "text": " And this time you're going to create one hot vector. So you're going to put zero, one, two, three."}, {"start": 1254.24, "end": 1258.24, "text": " So you want to put one here and all of the zeros will be here."}, {"start": 1258.24, "end": 1266.24, "text": " So this is just a different representation of this very same information. OK, so that's the encodings data structure."}, {"start": 1266.24, "end": 1272.24, "text": " And why do we do that is because we can then multiply this encodings data structure with the embedding table."}, {"start": 1272.24, "end": 1278.24, "text": " So, again, embedding table is the code codebook table. So you'll have just a second."}, {"start": 1278.24, "end": 1283.24, "text": " So we have K here. So this will be K dimensional and this will be 64."}, {"start": 1283.24, "end": 1287.24, "text": " So these are this is your codebook vector code table."}, {"start": 1287.24, "end": 1298.24, "text": " So if you can see here, just by multiplying this encodings data structure with the codebook vector, you'll end up extracting vector number three from from here."}, {"start": 1298.24, "end": 1305.24, "text": " And that's basically the quantization process we're talking about. So you'll end up having basically quantized vectors."}, {"start": 1305.24, "end": 1313.24, "text": " You'll have you'll end up having if you multiply these, you can see you end up with n times 64 dimensional table."}, {"start": 1313.24, "end": 1318.24, "text": " And those will contain the quantized vectors. So that's how the quantization is done."}, {"start": 1318.24, "end": 1328.24, "text": " Simple matrix multiplication. And finally, the interesting stuff here is so this E latent loss is this term here."}, {"start": 1328.24, "end": 1333.24, "text": " So that's this term here. As you can see, we put a stop gradient onto the codebook vectors."}, {"start": 1333.24, "end": 1339.24, "text": " And as you can see here, so these are the codebook vectors. We put a stop gradient, which is implemented as detached in PyTorch."}, {"start": 1339.24, "end": 1347.24, "text": " We do the same thing for this second term. So this is this part where we put stop gradient on the encoded vectors."}, {"start": 1347.24, "end": 1351.24, "text": " So encoded vectors are called inputs in our code. And we just do detach. And that's it."}, {"start": 1351.24, "end": 1355.24, "text": " Just simple MSC laws there. And again, this is beta. So that's your loss."}, {"start": 1355.24, "end": 1363.24, "text": " OK. And now for the funny part, and I'll just ignore this, is how do we actually implement these straight through gradient?"}, {"start": 1363.24, "end": 1367.24, "text": " So this is the way we implement that estimator."}, {"start": 1367.24, "end": 1375.24, "text": " So if you take a look at this equation, obviously inputs and inputs will kind of cancel out and you end up with quantized."}, {"start": 1375.24, "end": 1383.24, "text": " So the reason they added this part here is because when we do backprop, this part is basically excluded from the computational graph."}, {"start": 1383.24, "end": 1392.24, "text": " And so we end up with inputs. And so that means that basically whatever the gradients you get for your outputs, for your quantized variable here,"}, {"start": 1392.24, "end": 1397.24, "text": " those will just be copy pasted into the inputs. And that's the part I just explained pictorially here."}, {"start": 1397.24, "end": 1402.24, "text": " Basically, you have whatever the gradients here are. You're just going to copy paste them into here."}, {"start": 1402.24, "end": 1406.24, "text": " And once you have that, then you can just do because everything is differentiable in the encoder."}, {"start": 1406.24, "end": 1412.24, "text": " You can figure out the gradients for every single weight and then you can just update your network. OK."}, {"start": 1412.24, "end": 1418.24, "text": " So you may wonder why we are not using the KL term. Let me try and give you some intuition there."}, {"start": 1418.24, "end": 1430.24, "text": " So they say here, since we assume a uniform prior for Z, so compare that to VAEs, which had like isotropic Gaussians as the basically standard Gaussians as a prior."}, {"start": 1430.24, "end": 1435.24, "text": " Here we have we assume uniform prior for Z. The KL term that usually appears in the elbow."}, {"start": 1435.24, "end": 1442.24, "text": " So this is your evidence lower bound is constant with respect to the encoder parameters and thus can be ignored for training."}, {"start": 1442.24, "end": 1448.24, "text": " So let me try and digest this for you a little bit. So that's so pure."}, {"start": 1448.24, "end": 1455.24, "text": " If that if we take a look at that, we have a uniform distribution here. OK."}, {"start": 1455.24, "end": 1465.24, "text": " So so this will have K. Points in the domain here and what we'll end up with."}, {"start": 1465.24, "end": 1476.24, "text": " So how our posterior looks like, so how our approximate posterior Z given X looks like is it's as I mentioned, it's a deterministic function pretty much."}, {"start": 1476.24, "end": 1481.24, "text": " You will have you'll end up having a single vector that's the closest to your encoded vector."}, {"start": 1481.24, "end": 1488.24, "text": " So you'll end up with something like this. And if you take a look at what you get by doing a KL divergence between these two."}, {"start": 1488.24, "end": 1498.24, "text": " So if you do KL between these two, you'll end up with basically just one times log."}, {"start": 1498.24, "end": 1506.24, "text": " Then we have in the numerator, we'll have one and in the denominator we'll have we'll have because this is uniform, we have K points here."}, {"start": 1506.24, "end": 1515.24, "text": " This is going to be one over K probability. So the probability that corresponds to this one, let's maybe that's this one, will be one over K."}, {"start": 1515.24, "end": 1522.24, "text": " So you end up with KL divergence that's constant and that's equal to log of K."}, {"start": 1522.24, "end": 1530.24, "text": " OK. And so you basically can't minimize it because it's constant. And they mentioned it somewhere here."}, {"start": 1530.24, "end": 1539.24, "text": " So our proposal distribution blah, blah is deterministic. And by defining a simple uniform prior of Z, we obtain a KL divergence constant and equal to log of K."}, {"start": 1539.24, "end": 1547.24, "text": " So that's why they can't minimize it because it's constant in this setup. And yeah, I just wanted to explain that a little bit more."}, {"start": 1547.24, "end": 1556.24, "text": " OK. So for the prior part, they say here, whilst training the VQVAE, the prior is kept constant and uniform, as we just saw."}, {"start": 1556.24, "end": 1563.24, "text": " After training, we fit an autoregressive distribution over Z, P of Z, so that we can generate X via ancestral sampling."}, {"start": 1563.24, "end": 1572.24, "text": " So all of this sounds maybe a bit scary. Let me try and demystify this part. So if you're familiar with language modeling, if you watch my GPT-3 video,"}, {"start": 1572.24, "end": 1579.24, "text": " I'll link it somewhere here or like transformers, whatever, if you're familiar with language modeling in general, this will be really easy."}, {"start": 1579.24, "end": 1592.24, "text": " So what you basically do here is in a way you just try to predict the next token. So in the case of language modeling, you're trying to predict a token that represents some subword or something or like a word or a character."}, {"start": 1592.24, "end": 1597.24, "text": " Here, you're just trying to predict the next symbol. But like the semantics of the process is the same."}, {"start": 1597.24, "end": 1605.24, "text": " So imagine you flatten this into a sequence and you have the same, like basically same setup as for language modeling."}, {"start": 1605.24, "end": 1612.24, "text": " So let's say we have a simple example, not like a not like a matrix of symbols, but maybe just like a three symbols."}, {"start": 1612.24, "end": 1622.24, "text": " Like let's say we have three, two, five. OK. And how you train this all regressive model is the following. So you'd input something called some special symbol called start of a sequence."}, {"start": 1622.24, "end": 1628.24, "text": " Then you'd input three here and then you'd input two here. And as you can see, you're trying to predict."}, {"start": 1628.24, "end": 1635.24, "text": " So given the special token, you're trying to predict three. Then given three, you're trying to predict two because that's the next token, as you can see here."}, {"start": 1635.24, "end": 1642.24, "text": " And now given two, you're trying to predict five because that's the next token in the sequence, as you can see."}, {"start": 1642.24, "end": 1653.24, "text": " So this goes as an input to your model. These are the targets. And basically you're just you'll you'll basically use some like cross entropy and you'll backprop through the weights of the model."}, {"start": 1653.24, "end": 1657.24, "text": " And that's how you're going to train this this this token predictor model."}, {"start": 1657.24, "end": 1666.24, "text": " And it will going to be all regressive, which means once you output a token here, you're going to feed it back to the input and then you're going to keep on generating like that."}, {"start": 1666.24, "end": 1673.24, "text": " So once you have it trained on this discrete space that was so basically there are two things you do here."}, {"start": 1673.24, "end": 1680.24, "text": " First things first, you basically train this VQVAE end to end. So to learn how to reconstruct images."}, {"start": 1680.24, "end": 1688.24, "text": " And then you have you basically have a codebook that's that's fully trained and you have encoder and decoder, which are fully trained."}, {"start": 1688.24, "end": 1693.24, "text": " You just freeze all of those and you just train this predictor model."}, {"start": 1693.24, "end": 1703.24, "text": " And then when you want to generate new novel images, what you do is you just prompt the the all regressive model with the special token and it will just output some number."}, {"start": 1703.24, "end": 1710.24, "text": " Let's maybe say three and you feedback the three inside of the input of your all regressive model and it will generate something else like maybe four."}, {"start": 1710.24, "end": 1717.24, "text": " And you do that all the way until you have enough elements in your in your like latent space."}, {"start": 1717.24, "end": 1720.24, "text": " So that will be thirty two by thirty two in this example here."}, {"start": 1720.24, "end": 1725.24, "text": " And once you have that, you just feed that through this frozen decoder and out comes some novel image."}, {"start": 1725.24, "end": 1728.24, "text": " So that's the explanation of how this whole system works together."}, {"start": 1728.24, "end": 1731.24, "text": " Hopefully you you now understood it a bit better."}, {"start": 1731.24, "end": 1736.24, "text": " OK, so having demystified that, hopefully let's see some results."}, {"start": 1736.24, "end": 1741.24, "text": " And first thing they mentioned here is that they can compress their data by a lot."}, {"start": 1741.24, "end": 1754.24, "text": " And that's hopefully obvious. But like so they say here in this experiment, we show that we can model X, which is hundred twenty eight by hundred twenty eight in their example by compressing them to thirty two by thirty two discrete space with K equals five hundred twelve."}, {"start": 1754.24, "end": 1761.24, "text": " So case the size of that codebook codebook table, why are purely the convolutional PX given set."}, {"start": 1761.24, "end": 1773.24, "text": " So decoder. So a reduction of and you can see how much here and basically nine is here because if you have a naive encoding method of your symbols, you'll just have nine bits."}, {"start": 1773.24, "end": 1781.24, "text": " You need nine bits in order to represent five hundred twelve symbols if you assume that all of the code like sizes are fixed."}, {"start": 1781.24, "end": 1785.24, "text": " OK, and basically you can see here some reconstructions."}, {"start": 1785.24, "end": 1790.24, "text": " And what I want you to notice here is because this will be important for the later Dali paper."}, {"start": 1790.24, "end": 1794.24, "text": " I'll create like a paper over. I'll do and VQ again as well."}, {"start": 1794.24, "end": 1797.24, "text": " You can see that there is some blurriness to these images."}, {"start": 1797.24, "end": 1804.24, "text": " If you take this dog here and you'll see like the reconstruction is pretty good, but there is some blurriness happening here."}, {"start": 1804.24, "end": 1806.24, "text": " It's especially noticeable here."}, {"start": 1806.24, "end": 1811.24, "text": " And for this real image here, you can see it's much blurrier."}, {"start": 1811.24, "end": 1812.24, "text": " You can take a look at the paper yourself."}, {"start": 1812.24, "end": 1815.24, "text": " But that's that's something I want you to notice here."}, {"start": 1815.24, "end": 1817.24, "text": " So those were the reconstructions."}, {"start": 1817.24, "end": 1829.24, "text": " And here after we do that sampling process in the inside of the latent space using that order aggressive model, you can sample novel unconditional images and you can see some results here."}, {"start": 1829.24, "end": 1831.24, "text": " And you can see it's really diverse."}, {"start": 1831.24, "end": 1835.24, "text": " And that's one of the advantages these likelihood models have compared to GANs."}, {"start": 1835.24, "end": 1850.24, "text": " They cover all of the modes because of the way they are trained, whereas GANs sometimes can omit some modes or they can even like go into something called mode collapse where they basically output just a single image or similar looking imagery."}, {"start": 1850.24, "end": 1859.24, "text": " OK, they also do this not only for images, they can also compress and reconstruct audio, video, et cetera."}, {"start": 1859.24, "end": 1861.24, "text": " So I found the audio part really interesting."}, {"start": 1861.24, "end": 1865.24, "text": " So this is your original audio after you reconstructed."}, {"start": 1865.24, "end": 1867.24, "text": " You can see the waveforms are not quite the same."}, {"start": 1867.24, "end": 1870.24, "text": " But what's interesting is that they have the same content."}, {"start": 1870.24, "end": 1877.24, "text": " Couple of nice remarks they made here is it can be heard from the provided samples and as shown in Figure 7 reconstruction."}, {"start": 1877.24, "end": 1879.24, "text": " So Figure 7 is this one here."}, {"start": 1879.24, "end": 1881.24, "text": " You can find the samples on the website."}, {"start": 1881.24, "end": 1882.24, "text": " You can hear them at your own pace."}, {"start": 1882.24, "end": 1888.24, "text": " But like the reconstruction has the same content, same text contents, but the waveform is quite different and prosody."}, {"start": 1888.24, "end": 1892.24, "text": " So the way you accent your sentence in the voice is altered."}, {"start": 1892.24, "end": 1908.24, "text": " So this means that the VQVAE has without any form of linguistic supervision, so without any supervision, learned a high level abstract space that is invariant to low level features and only encodes the content of the speech, which is awesome since we don't have any labels here."}, {"start": 1908.24, "end": 1917.24, "text": " So while samples drawn from even the best speech models like the original WaveNet sound like babbling, samples from VQVAE contain clear words and part sentences."}, {"start": 1917.24, "end": 1931.24, "text": " So what I recommend you do is you go ahead and listen to the samples from this model and from WaveNet, which was very popular, like parallel WaveNet, like that was a full up work to WaveNet, was actually used in production at Google."}, {"start": 1931.24, "end": 1932.24, "text": " It's maybe still used."}, {"start": 1932.24, "end": 1935.24, "text": " I'm not sure, but I know it was used in production for Google Assistant."}, {"start": 1935.24, "end": 1939.24, "text": " OK, so I think that's pretty much it for this paper."}, {"start": 1939.24, "end": 1948.24, "text": " Let me just briefly go and explain you how this VQVAE 2 version differs from this version 1."}, {"start": 1948.24, "end": 1952.24, "text": " So first thing you can notice is it's got much better reconstructions."}, {"start": 1952.24, "end": 1955.24, "text": " You can generate much better imagery using this model."}, {"start": 1955.24, "end": 1958.24, "text": " The main difference is so everything remains the same."}, {"start": 1958.24, "end": 1960.24, "text": " The loss remains the same."}, {"start": 1960.24, "end": 1966.24, "text": " The only difference is they have this hierarchical structure of the of the latents as well as the priors."}, {"start": 1966.24, "end": 1976.24, "text": " So what I do instead is they have this top level codebook vectors and they have this mid level codebook vectors and they are of different resolutions."}, {"start": 1976.24, "end": 1981.24, "text": " So that means this one will have to focus on some like higher level structure of the image."}, {"start": 1981.24, "end": 1984.24, "text": " This one can maybe even focus on some details, texture, et cetera."}, {"start": 1984.24, "end": 1989.24, "text": " And once you combine them into decoder, you can reconstruct much higher quality imagery."}, {"start": 1989.24, "end": 1990.24, "text": " Same thing here."}, {"start": 1990.24, "end": 2004.24, "text": " You train this autoregressive model in this latent space and then you use that to condition the next level prior autoregressive model to basically regress this new latents and then you can decode them into imagery."}, {"start": 2004.24, "end": 2007.24, "text": " So that's the only difference pretty much between this model and the last one."}, {"start": 2007.24, "end": 2009.24, "text": " So they're using this hierarchical structures."}, {"start": 2009.24, "end": 2012.24, "text": " And as you can see, they get much better imagery."}, {"start": 2012.24, "end": 2016.24, "text": " Like this is even compared like comparing even this to, I don't know, style again."}, {"start": 2016.24, "end": 2023.24, "text": " V2, the images look really, really cool, although style again can can generate much higher resolution imagery."}, {"start": 2023.24, "end": 2024.24, "text": " I think so."}, {"start": 2024.24, "end": 2028.24, "text": " But yeah, either way, this is really cool for a VAE."}, {"start": 2028.24, "end": 2033.24, "text": " Usually GANs have much better reconstruct image images than VAEs."}, {"start": 2033.24, "end": 2036.24, "text": " OK, here are some more images."}, {"start": 2036.24, "end": 2039.24, "text": " Again, I think this is class conditional on ImageNet."}, {"start": 2039.24, "end": 2041.24, "text": " I think doesn't matter that much."}, {"start": 2041.24, "end": 2042.24, "text": " Let me just see."}, {"start": 2042.24, "end": 2044.24, "text": " Yeah, these are some high resolution images."}, {"start": 2044.24, "end": 2046.24, "text": " You can see they're pretty good."}, {"start": 2046.24, "end": 2048.24, "text": " You can see it's still not perfect."}, {"start": 2048.24, "end": 2050.24, "text": " There is some distortions here happening."}, {"start": 2050.24, "end": 2051.24, "text": " I don't know. Maybe I'm wrong."}, {"start": 2051.24, "end": 2053.24, "text": " Hopefully you found this video insightful."}, {"start": 2053.24, "end": 2058.24, "text": " Let me know what you think down in the comments, whether you like this format, whether the code part was useful."}, {"start": 2058.24, "end": 2063.24, "text": " If it was, I'm going to try and include more of these into my next videos."}, {"start": 2063.24, "end": 2077.24, "text": " So until next time, subscribe and share this video and bye bye."}]