JAX Diffusers Community Sprint Talks: Day 2

Key Takeaways

hello everyone and welcome to the second day of Jax diffusers um Jacks diffusers kickoff events and today we are going to host the team Suraj and Sabrina and without further Ado I would like to have Tim team is working in Google brain and today he is going to talk about efficient image and video generation with distilled diffusion models the floor is yours thanks thanks for the intro right so for today I will actually talk about basically six different papers that collaborators and I have written over the past two years so this is Joint work with lots of great colleagues at Google and I'm very grateful to them for all their amazing work of course it's quite a little bit of material so I'll just give a very high level overview hopefully you'll know where to find the good bits after this talk and of course I can answer more detailed questions also after [Music] to set the stage and reflect upon what an amazing year has really been for generative modeling so I've been in this field for a while now and I think the past year has really been here generative modeling so at Google we released imagine an imagining video two models that I was involved in but also for Nike is another epigl model that we released and a party a text to image model and of course outside of Google we had tally 2 and stable diffusion all basically within one year um and we really came quite far in the years before so when I got started in these fields in 2013 2014 we just had feas be invented and Gans and what it barely works if you look at the samples on the left here you can't really make out what we're generating um so then we made some advances in Auto regressive models MSL was at open AI at the time and I had this opportunity to present some of my work to Elon Musk and I actually showed them these particular samples right here well needless to say he was not very impressed um but fortunately we made a lot of progress as a field since then in 2019 I think we made a big step with big can where for the first time we were really able to synthesize compelling images still in a very limited data set imagenet where we had specific classes and corresponding images um and then again but then in the past year really is for the first time we're able to take basically arbitrary tax descriptions of what we want to generate and then generate truly compelling images at high resolution um so for Imagine specifically I want to tell you a bit more about the system the way it works and what it can do um it's at a high level it's a text to image model when you put in a string like sprouts and shape or text imagine coming out of a fairy tale book and out comes a 1K or 1K pixel a high resolution image with quite a bit of intricate detail as you can see here um it can do amazing things like spelling and also composition so here we're asking for a transparent sculpture of a duck made out of claws in front of a painting of a landscape and it actually gets all these corrects and knows how to compose those into a coordinate image and it can even do this when you put in things that it probably has never seen in training at least not in this combination like a photo of a rainbow called a cow floating in the space above Earth milk away in the background um it can even generate objects that cannot exist in the Real Worlds like this coffee cup or it can generate objects that we would not like to exist in real world because they would be very flexible um so there's a number of things that go into imagine that make it possible and there's four main parts there's the data so the model itself is trained on close to a billion text image pairs so really quite a large scale the model itself is based on diffusion models of course but we combine them into something we call a Cascades so it's a combination of different models it also relies heavily on Frozen text models so large language models um and we there's also some advances in how we sample from this model so we use classifier free guidance and we have a number of tricks to maximize the text alignments between the image we generate and the text prompts and then of course we apply this at quite a large scale which is necessary to get the best possible results so I'll say a little bit more about these three aspects cascading the text models and the classifier free guidance so at the basis the diffusion model that we all know and love it turns Pure Noise into an image but we compose it in a pipeline so the first stage of the pipeline we take the text prompt we embed it into effective space using a frozen language model so it's not trained and get out a set of embeddings um the first diffusion model then takes the embeddings and some input noise and turns it into a low resolution image the second diffusion model takes the embeddings and the lower resolution image and up samples it to a higher resolution image and we do this once more to get to our final results and there's some tricks to make this work for example you need to have some augmentation whenever you train these super resolution models so that they are robust enough to actually work at test time and we have a number of details on that in the paper and taking all taking it all together this is a very effective way of generating high resolution images of the highest quality um one other very important trick that we use in this model is called guidance so in the diffusion model what we do is we have noisy data which I denote by XT here and maybe a class label or a text bronze why so in this case y might be the label of dog and this gets and then we feed this to a model that tries to denose this to the clean image that I denote by X naught here um however in practice what we often do is we parameterize the model in such a way that we actually predict the noise rather than the image itself so this is can be thought of as a residual parameterization and this noise estimate actually turns out to be proportional to the gradient of the log density of the image given the textual description or the label um openai a whole bunch of folks at open AI in 2021 then came up with this this great trick to modify the prediction of the Noise by adding in the gradient of a classifier model so the classifier model takes in the noisy image predicts a dog or some other class and then what you do is you move the noisy data in the direction of the class that you want to generate so this ensures that you're generating something that's well aligned with the text prompt um we did some follow-up work on this using a technique that we call classifier free guidance where instead of an explicit classifier or use an implicit classifier using base rule so we say the probability of the label given the data is proportional to the probability of the data given the label divided by the unconditional probability of the data so if we take that implicit classifier we take the log gradient again and then we get this other term which is the difference between the gradients with respect to X of the log density given the label and the gradient of the log density um in conditional of the data of the of the prompt and this difference you add into the estimate of the noise and this works quite well um in practice we predict noise so when you plug in the noise prediction rather than the gradient estimates that looks like this so we have a noise prediction conditional on the text prompt and one in conditional and in practice we do do that by simply training one conditional model that always takes in embedding for the text prompts but sometimes we set this to zero so we just encode the zero string and this we take to be the unconditional prediction of the image or the noise so this trick works quite well in ensuring that what we generate corresponds to the text um so here I'm showing you the results for a number of different settings for this uh this W parameter and guidance on the y-axis I have FID so FID is a matter a measure of how well our model is capturing the data distribution and on the x-axis I have the clip score so clip scores a measure for how well what we generate actually represents the text prompt that we asked for so as we increase the guidance weight we move from the bottom left of this image to the top right so initially both FID and clip score improve and then as we increase guidance further FID starts to deteriorate so what we generate is further removed from the data distribution then for the null guidance case but the alignment is still improving and what we find in practice is that users often care more about the alignments than about the Fidelity with which we capture the underlying data distribution so when we combine this with some tricks where we do some clipping or thresholding to make the model stable this this results in a red line here where we can get very high clip scores and this corresponds to Generations that or images that the users find very compelling um for the text model we investigate different classes of text encoders we look at language models like D5 and birds and we also look at texting encoders from image to text models like clip and we actually find that these tax only models they actually work very well which is kind of surprising since they never even saw visual data during training um so we take these language models and we freeze them so we don't train their parameters and we still find that their representation actually is very useful for generating images and in particular what we find is that larger text models give us representations of texts that are more useful or that the model finds easier to use in generating images and we actually find that using the largest text models actually makes a large difference for the perceptual quality of the images that we can generate um so for us after having done Imagine The Next Step kind of logical Next Step was to ask whether we could do the same for video and it turns out what we can so we have this follow-up project called imagine video and where we generate short five second clips of video um in much the same way as we do text conditions images in Imagine um like I imagine for images imagine video is a really a pipeline of different models so we once again start with a text prompt that we feed into a large language model D5 XXL in this case so that's the yellow box that you see here after that we think it's embedding and we feed it into the first diffusion model which is these in the blue box it's a a base model that generates a very low resolution video so in this case just 16 frames at a resolution of 40 by 24. um so for about five seconds that's at three frames per seconds after that we then up sample this low resolution video in multiple stages so in all the red boxes that you see in this graph we do temporal super resolution meaning that we increase the frame counts of the video and on all the green boxes what we do is spatial super resolution so there we increase the um the pixel counts or the the resolution of the the frames of the video so we do this six times to arrive at our final output which is five seconds of 24 frames per second video at high resolution um so I'll give you some examples of what it looks like later the mole itself at each of these stages is really quite a simple diffusion model in a previous paper we introduced this model that we call a video diffusion model and it's really it's kind of the the the smallest Delta we can make from the image model um so basically what we're saying is that a video is is kind of like an image you just have one additional dimension in the tensor and similarly what we do is we take the standard to the unit that we used for images and we just apply it to each of the frames in a video and we add some additional layers that exchange information across time so it's a standard unit with some additional attention layers that exchange information between the frames so here's one example of something we can generate with this and here we ask for a bunch a bunch of autumn leaves falling on the lake to form the text imagine video and quite amazingly it can both do the spelling and it can also get the timing right so here's another example so what I find quite amazing about these examples is that it comes with the spelling but it also gets the timing exactly right so the text actually starts appearing at the start of the video and completely finishes uh appearing at the end of the video and this is actually not really Cherry picks we can quite generate this quite consistently um although there's nothing in a model that actually tells it that it needs to take exactly five seconds to generate this and there's nothing in the training data either um so really the magic that actually allows this to happen is guidance um so in guidance we ask the model to generate something that aligns well with the text prompt within the span of time that we give it and the only way it can do this is by actually matching up the text appearing with the duration of the video and this is also quite quite unique to this wave generating where I regenerate all the frames at the same time if you would do something also regressive where you would generate one frame at a time it's much more difficult to get the model to actually understand this timing without explicitly telling it that the video is five seconds or that the C needs to be completed by a certain number of frames and here's some more examples I can do things that are out of distribution like a wooden figurine surfing on the surfboard in space um another out of distribution example on a happy elephant wearing a birthday hat walking under the seat and here's a very cinematic one flying through an intense battle between Pirates ships in a stormy ocean and the last one I feel for Castle so um it's actually great at doing these aerial shots and there's a lot of Drone footage in the data I think and yeah I find these very nice also um of course we also evaluate our models on a number of standard benchmarks and especially for the base video diffusion model um so here we look at ucf101 which is a benchmark for unconditional generational video and there we find that we get a very very good results in terms of both FID and Inception score so here FID is um again captures the how accurately we represent the data distribution and Inception scores more related to alignment similar to the clip score I showed you earlier and on both we do quite well and we can also do video prediction so in video prediction we actually get the first frame of a video or the first couple frames and the task is we need to play out a video from the starting point so predict where the scene is gonna go how things are going to develop and the most well known Benchmark here is this rubber pushing tile Square you get lots of videos of this robot arm moving around and your task is to predict how it will interact with the objects and what will happen um so there we do quite quite well um another video prediction task is kinetics so kinetics is a collection of videos like you might encounter on YouTube um of people performing various actions and also there with the predictions work quite well especially if we use more fancy sampling methods like lunge event diffusion I'll talk about I'll say a bit more about that later um so I've showed you what these models can do in practice or at least if you're willing to take some time and maybe do some cherry picking and try out a number of things but it doesn't mean that this really makes for a very compelling user experience because naively sampling a video currently takes about 10 minutes slightly more for a five second video which means that if you're interacting with this system uh you know it gets boring very quickly so I strongly believe that generated models need to be very fast in order to be practically useful you have limited patients as a user you want the system to be interactive and so the systems need to be fast our Solutions we make these models faster is to distill diffusion models into models that take much fewer sampling steps and we typically go down to four to eight sampling steps without losing much quality um we developed this technique in two papers Progressive distillation were fast sensing of diffusion models so I clear paper and a more recent cvpr paper and that's on archive for reading and will be featured at the next cvpr and what we find is that if you apply these techniques to the imagining video samples that I showed you you can actually generate these in about 30 seconds which I think is just about the amount of time that is that you're still willing to wait and where you can still try out different problems and interact with the system we also have some open source codes that demonstrates the the basics Behind These algorithms it's all written in Jax so it should play nicely with all the other stuff that you're building and experimenting with so before I explain how exactly we do this distillation I want to say a little bit more about how we actually sample from diffusion models in the first place so here on the left you see the standard view which is that sampling is done by assimilating a stochastic interferential equation so we get this this noise from a known starting distribution like a gaussian and then we simulate this stochastic process that gradually takes us into the direction of the data which is much more as a much more complicated distribution so what does that looks like is we gradually turn noise into the clean image and we do that by simulating the stochastic differential equation at the bottom here but it turns out that actually there's an equivalent view where we have very smooth and deterministic paths from the initial random noise to the data distribution and these parts are given by something that's called the probability flow ode or the ordinary differential equation which is given at the bottom here so here there's no noise term it's a deterministic process that we can actually compute exactly if we have enough time turns out you can rewrite the stochastic differential equation to make it more similar to the ode so here I've split out in the green boxes the deterministic parts and then the remaining parts is language event Dynamics so what this basically says is that both the stochastic way of sampling from these models and the deterministic way are basically the same thing the only thing you get in addition when you do the stochastic sampling is some additional mixing um according to the data distribution using an mcmc algorithm like language of an Dynamics so it's this green box that we can actually help to distill so the green box is something you can integrate for a number of steps and then it gives you a deterministic mapping from one point of the the input noise to something along the direction to the the outputs image and this deterministic mapping is something that we can hope to represent with the neural net that takes fewer steps than the sulfur that we use to actually generate the path um so that's what we actually distill um so in order to do this efficiently we propose an algorithm called Progressive distillation where we apply this distillation in multiple stages and at each stage we take a student model that is warm started from a teacher model just by making a copy and it learns to imitate the teacher model but it actually uses half as many sampling steps so what that looks like is something like this on the left here on the left here you see a teacher model that takes four steps and what we do is each of these two for two steps we try and distill that into a student model that does the same thing but with only one step so the first two steps get distilled into something that takes one step to get to the same point the second two steps get distilled into something that also takes one step and now at the next stage we take this new student model and we turn it into the teacher model we repeat this for this process we now go from two steps into one step and so you can imagine if you start out this way with many steps you can iteratively half the number of steps required and this works quite well in practice and it's very computationally efficient so looking at the algorithm it looks like this so on the left is the standard diffusion training codes and on the right is Progressive distillation and the main changes are highlighted in green so I'll go through them step by step so in Progressive distillation we start out with a trained teacher model so the way I've written it here is it's a prediction of the clean data X and with it using as input and noisy data as a set t at each iteration of progressive distillation we do a warm start so we say we set the students equal to the teacher by copying the parameters and then comes the interesting part so the interesting part is where we set the targets that we train the model against and we determine that Target by running two steps of the ddim integrator so ddim is one way of integrating this ode that I showed you in the past slides um and it's the most efficient integration that we have so we use two steps on ddim to get you a next step in the sampling chain um then from that next step in the sampling process we back out the the targets that are our student model we need to predict in order to get to the same points in one step so that's given by this formula and then this target is what we actually plug in into the the loss for the model so this is what we ask the student model to predict so we do this for a number of steps until we're happy with the results and then the student becomes the teacher for the next iteration of this algorithm and we halve the number of sampling steps so what we find is that this works quite well in certain data sets and what I'm showing you here is the FID for a distilled model for varying number of sampling steps and I can I'm comparing it to ddim and to a stochastic Baseline sampler drive that is separately optimized for each number of sampling steps so it's quite a tough Baseline to beat and we find that as you decrease the number of sampling steps by a lot and they still model does much better and in particular we find that four to eight sampling steps is kind of the sweet spots where you don't lose Mars in terms of quality um but you do have a very large speed up compared to naively taking as many as 5-12 steps in this case so that was the first paper where we showed this as a proof of concept on image stats and on other data sets it is more recent cvpr paper we extend this in multiple ways so we showed that this approach also works with classifier free guidance we show that you can actually use the distilled model to do stochastic sampling not just deterministic sampling um we show a little words for taxi image models not just class conditional generation and we do image image translation we're in painting and we show that this actually also works for latent diffusion models and so most of the samples that you see in these slides are actually generated using a distilled model or stable distilled version of stable diffusion um I'll tell you a little bit more about some of these aspects um so in particular and the way we actually distill models with classifier free guidance is in two steps in the first step we take this classifier free guidance model that actually takes two Parsons in order to take one step in the sampling algorithm and we distill that into a model that only takes a single pause we do that by sampling a Gardens coefficient so that's the the W parameter I showed you earlier and we sample this from a particular range like between zero and four then we find the teacher prediction that corresponds to that so we did two passes through the the teacher model we apply the guidance formula and we get a prediction out and then we try and mess up with the student model and the student model takes us input also the guidance weight w so in addition to the input data and embedding of the time step it also takes an embedding of the W parameter so you train that for some time after which you have a single model that now does in a single step what the old model did in two parcels and after that you can just do the progressive distillation that I previously explained and so you're at the number of steps that you're happy with um now we have a deterministic model that takes many steps and previously we only did deterministic sampling using that approach so that corresponds to the top of this uh this slide where in this case we take for deterministic steps I use it shown by the blue line to get to the image and then we give us output and it actually turns out that we can also do stochastic sampling with this and the way that works is that we just simply double the step size that we use for all these domestic steps so the blue steps are twice as big but then after each determinist accept forwards we actually take a step backward by adding noise back into the image so adding those back into the image is just exactly the same as what we do during training we're taking a clean image or a halfway clean image you add a little bit of noise in um and because we can actually split out this sde into a deterministic part and a a noisy part um you can actually show that this is a very this is actually also a very uh Accurate Way of sampling from this process and what we find that this is particularly useful for cases where we use high guidance weights um lots of lots of numbers here I'll show you two in this red box here so in this red box on the left you see the number for the deterministic sampler the Inception score that you achieve with a deterministic sampler using eight sampling steps and then guidance way to four and the right number in the red box is the same number you get by doing stochastic sampling we find this is generally a bit higher when using uh when using a few steps in high guidance weights um and so this uh and this also translates to images that we find a bit more compelling um yeah finally what we find is that this is directly this algorithm is directly applicable to latent diffusion models like stable diffusion um so uh channeling the first author of the paper apply this to a stable diffusion where you can get quite good results using as few as two sampling steps so to summarize uh the talk so far um I think generative models are really no longer just a science projects when I got started in this field I clearly were um I mean no one is going to use the samples on the left here but now we're at the points where these models actually produce quite amazing things um and I'm I'm confident so they will end up in lots of products in the coming years and will be very impactful diffusion is becoming quickly the new paradigm um for achieving state-of-the-art results using generative models um it's a very flexible framework it's it's stable it's not like again where we don't really have a proven convergence it's a that's more of an empirical science but really we have a simple loss term that we can and that we can minimize and use the trainer model and finally I believe strongly that these moles need to be fast if we actually are going to use them and um our approach to that has been distillation and that's been proven very effective so far right thanks I can now take some questions foreign this was my first time seeing distillation outside of Transformers honestly I really like the idea and the intuition behind it and we have a bunch of questions so Patrick is asking and not sure how much Tim can talk about this but I'm very curious if Imogen would have similar performance if it would would have been trained only on public data um right so data is of course very important um there's a lot of data out there on the internet I'm sure you could get similar Performance training just on public data but then the question is how do you curate the data um so just taking every random jpeg you can find on the internet that won't work you need to put some time into actually selecting the right ones and if you do I'm sure it can work and um you know at Google we're definitely exploring lots of different data sets and finding that different data sets are appropriate for different purposes um but that yeah good results can be achieved in every variety of ways and the second question is can we do the same with all types of noise or distortions not just gaussian also is the noise required to be uniform over the entire image or can we have areas with more noise than others right that's a good question so there are some words in literature where they do use different types of noise but in general all of our results do depends on this infinite divisibility of the noise so one thing that's very unique about the gaussian is that if you add two gaussians together you still get a gaussian out or a quick the other way around if I give you the noisy data you can actually decompose this into a number of additive noise terms at all that are actually consistent with the output distribution so there's ways of doing something similar with different distributions but it's um it's it's quite tricky to get something that's theoretically Justified so I would say in general yes it needs to be gaussian and what we do find is that you can actually add more noise in in certain directions than in others and we actually have a paper called blurring diffusion models where we do a transformation of the image very effectively we're adding noise in particular directions more than it matters and this also works quite well yeah I will Emil talked about it a little yesterday because people were asking about it and our third question is from Patrick people are playing around with using character-based text encoders now such as byt5 or K9 uh hoping that they are better at generating exotic text what do you think about this right so so what we do find is that if you wanted to be able to do spelling like in this uh this last example that I'm showing you then of course you do need to actually represent the individual letters in in your embedding um so big Tech models like T5 already do that quite well but I I imagine that doing character based encoding actually is even better if that's what you care about um yeah there might be a trade-off the trade-off though between um capturing the feel of an image and exactly representing a prompt so something like a clip-based encoding like in Dolly 2 and also in some other things that we tried um that is actually quite effective for capturing the general feel of a prompt so in italy2 paper there's some very nice examples of that where they take a logo they encode it and they generate multiple logos that's kind of the same um and that's something that I'm not quite sure you would get from an embedding from a text mobile that has never seen visual data another question is this sequential approach with many large blocks seem quite expensive they didn't notice much poorer inference times for video Imogen um so as I mentioned for for Imagine for for video imagine for the the base model where we take seven different steps seven different diffusion models um and we use 256 to 128 steps to to sample from them it does take over 10 minutes to generate a short five second video so that is indeed quite slow so there's different ways of speeding that up um one approach to the distillation that I talked about here another is the the one that Emil photo out yesterday and where rather than having a sequence of multi-resolution models we're actually trying to generate at the highest resolution from the start we can actually do that quite well now mainly by making some changes to the noise schedule that we use for sampling um and this is definitely something that we're still working on and I think these modes will get faster and more efficient um as we make progress there another question is at what video length assuming the memory allows for it to do the video videos become inconsistent right right um that's a good question so in a video diffusion models paper we also present a way of Auto aggressively extending this video um so there's advantages to generating the whole video in one go which is what we do in the Imagine video but then if later on after training for five second videos you decide to actually want them to extend the extended you can actually do that using these models um so this is a proven concept we actually try doing this and generating like minutes long video for for a particular text prompt and it does actually work quite well we don't we don't really ever diverge um it's just the alignment between the text prompt and the video can deteriorate a little bit so for example if you're generating a person lifting an item from the table then the first five seconds with high guidance that's exactly what happens then if you extend that video you know you see like a person with a object in their hands like standing around a little bit and it's not that interesting um so you don't think I've gained that much from it um but it's it's certainly possible and it doesn't diverge and that's that's quite a nice property of these models and we can have maybe two more questions let's see um do you I think this is a good question do you imagine that Progressive distillation could be adapted to the distillation of models in other fields right right um yeah so that's a very General open question um so these diffusionals though have a very specific structure where they take many steps to generate something that we actually know deterministically um and that that allows you to then amortize taking all the small steps into something that takes one large step so amortization is something that is used extensively in machine learning so for example in alphago you might take many steps to arrive at a particular move on a go board or or on a chessboard or something similar and then you can train your neural Nets to cast that move from the start without going through this whole search process so I think intuitively that's actually very similar to what we're doing here um and I'm sure you know that the basic concept is applied to many different uh applications in machine learning um in in another Fields uh and and yeah hopefully Progressive distillation can be an inspiration for for trying even more applications I would imagine and for uh stuff like alphago it would be quite expensive to get somewhere though yeah I mean that but that's true right like in an off ago also the computers in search or in every search that we evaluates in the neural Nets and if you search very deeply um then the the base model that you or the first time you evaluate the model that's only very small it's not part of the compute but then what they find in their case is that if you train is well enough then even the first step is pretty decent and this is similar to how our distilled model is still very slightly worse than the base model but it gets you most of the way um so let me see if there is like a cool question there is like a lot of cool questions but I'm trying to pick the best one because we will host surai soon um you say that diffusion become the dominant Paradigm and I agree do you think the recent consistency models could become the nail Paradigm One Two Step without distilling right so the the recent consistency models paper I quite liked um like mentally I still place this into the diffusion Paradigm it's uh very closely related and I actually think that there are strongest results roles are still distilling from a trained diffusion model so their Paradigm does allow for training from scratch but the best results that he achieved were still using distillation and the Boris is asking is a pipeline of multiple models necessary or could you use a large model end to end so using a large model end-to-end is something that we've explored and that's Emil talked about yesterday so I'm sure you can still see that talk also um it is a bit more tricky to to train as correctly but we're making a lot of progress there um yeah once again you need to adjust the noise distribution and maybe also adjust the model um but I believe we can indeed train one big model and this was the last question I really liked your presentation thank you so much for accepting to speak today thanks for the opportunity thank you see ya bye and now I will be hosting Suraj from hugging face and he will be talking about mask generating models which was something that I was very curious about would you like to share your screen yes okay can you see it um yes let's go okay I will leave you with your presentation good luck thanks Mary all right hi everyone I'm Suraj I'm a machine learning engineer at chugging place and I mostly do open source stuff all the commentator of Transformers and then now I'm a commentary of defeaters and I'm mostly working on open source XP image models so today I'm going to talk about this new method for image generation using discrete tokens to which is like a which was recently proposed by Google and it's called mask generative models so basically it's a Transformer based Master language model that can do really good image generation and the papers like musket and Muse have established state-of-the-art results on some image generation tasks at the time of publication so this talking I mean I'll briefly talk about why do we need these models and what are the limitations of current models and then I'll briefly present what Muse is and then we'll go deep into how these models are trained and how to do inference with these models So like um I mean as like everyone knows that the most dominant approach for image generation right now is diffusion models we really don't need any introduction for it and another really popular approach before diffusion was Auto regressive Transformer models you know like the first Dolly model and then Dolly mini and party which he uses like a two-stage approach where there's a VQ gun which can take an image you can create a discrete representation of it and then we use a Transformer model and encoder decoder Transformer model that can Auto regressively generate the image tokens by doing cross attention or the text and especially if party has like really really great results on these models but uh one of the limitations of both of these approaches is that these models are slow at sampling time okay one side note here that this the diffusion field is very fast moving and then we we have some new methods like consistency models and distillation which can bring down the sampling Time by by a lot but right now I say for example if we take stable diffusion we need to do like at least 20 or 25 steps to get good Generations if you have Auto regressive model like Dolly mini or party and if you when you're generating 256 tokens then the model you need to run the decoder for 256 forward passes so it's like the the inference cost increases by a lot and especially if you want to do high resolution generation or video generation then that becomes a bottleneck so instead these masking models like musket and Muse offer like a more efficient approach compared to this these other methods they use a bi-directional Transformers and these models because of the way they are trained they are able to predict multiple tokens simultaneously like like predict multiple tokens in a single step and that way like they can accelerate the decoding so instead of having to do like 256 steps you can do like four steps that steps or 16 steps depending on the kind of quality you want from the image so that's kind of a more motivation for these models so now we'll go over like how mask each actually works so musket stands for like mask generate you image Transformer and like the first Dolly and the party model it's a it's a two-stage model so the first stage is a VQ model I won't go into the details of VQ models here I think you can look into taming Transformers or there are lots of cool resources for understanding VQ models but basically what it does is like the VQ part also has an encoder and decoder the encoder takes an image and then creates a discrete representation of it so like you get like discrete integer tokens and the decoder can take those discrete tokens and turn them back into the image and then there's another Transformer model which is a bird like mask Transformer model and the way we train this model is we we get these discrete representations and then we mask a bunch of tokens in the um from the input for example here the gray tokens are mask tokens and then we feed it into the Transformer and then the Transformer is strange to predict these mask tokens let's get two stage approach say this is the VQ stage and there's the decoder stage this but the decoder is not Auto regressive it's it's a masked language model but it's a must model it's not so language so during training uh musket learns to predict randomly masked tokens by attending two tokens in all directions like this is one of the key thing here for example image images kind of have bi-directional tokens it's not like language where the where the future tokens only depend on the past past tokens and the past tokens are not influenced by Future token but like in images like you you have bi-directional dependencies is like instead of using a causal musket you just since it's a mask model it uses a bi-directional mask so each token can attend to every other token so it in a way it has more context and then at inference time uh we start with all masks tokens and then the model keeps on predicting the mask tokens and then we keep on iteratively improving the improving the generation that's kind of the basic overview of musket now uh during training what we do is as I said it's a two-stage model we use a VQ model encoder to get a discrete representation of token and then we sample a subset of tokens and replace them with a marks token this is the key part here um for example if you are doing Mass language modeling in NLP for so there's bird you use like a fixed masking rate of 15 and you mask like 15 percent of any random tokens but in in musk each or in Muse what we do is we use a scheduling function and that scheduling function tells us how many tokens to mask so the masking we do here is a variable we don't like we don't have a fixed masking rate of 15 percent because to this is to be able to like to emulate the inference scenario we use a variable masking red so the model can see like what to do when all to concern Mass what to do and like some partial tokens are masked and which which is a really key thing to get this model to work to use a variable masking rate and then the model is trying to predict these miles tokens by attending to all the tokens and Martin this like apart from the like the scheduling function the musket training is exactly similar to how you would train a bird model so but instead of the language tokens you have the discrete tokens of image given by the week you can you use a variable masking schedule for example for one example you would mask 40 in another example you would mask all of the tokens for some examples you would mask like just five percent ten percent on one percent of the tokens and then the model is trained to predict these mask tokens and then during decoding as you can see in the image what we do is we kind of fix the number of steps for example we say like we are going to do the inference in seven eight steps or 16 steps and then when we we are starting the generation at at t 0 we start with all mask tokens then feed that input into the model and then we then the model predicts those mask tokens and then we choose some tokens which have like a bit of a higher probability for example you check this like the dark gray area here is the token selected at that time step and then we keep on iteratively doing this like at each iteration we select some tokens which have a which have a high probability and then we mask out the remaining tokens and keep on doing this generation so at the last steps when we are here we then once we have like a lot more tokens generated we just at the last step we predict like all the remaining tokens and then we have our final image also the key thing you notice here is that when we start the generation the number of predicted tokens are very less for example in the first step there's just one token then here like seven or eight tokens but as we get more context we choose like multiple tokens which sample multiple tokens that's because the model has more bi-directional context and to to introduce the diversity in generation we use temperature and healing so like instead of choosing the token with the highest probability we just we say we set a certain threshold and depending on the like from that threshold we just randomly sample some tokens and this helps the model to like retain diversity in in regeneration and um yeah so that's pretty much how training and generation with uh must get works and another cool thing apart from these models being fast is that these models due to their nature of training support a lot of image editing tasks out of the box for example they can do zero shot in painting or outputting they can do image extrapolation and zero shot mask free editing like because the the model is trying to predict mask tokens you can to do in painting you can just like mask certain only mask out the tokens in The Mask region of the image and then ask the model to predict those and it it works pretty good as you can see in these examples again the key thing here is that these models are not trained to do in in painting out painting or like mask pre-editing it's just due to the nature of the nature of training of these models they can like zero shot they have the zero shot ability to do in printing Out Printing and mask free editing and the next step of musket is Muse which is basically like a musket model but it's for text to image conditioning so it has a text encoder uh as a conditioning and then the model text like as musket it takes the mask tokens but also does cross attention over the text embeddings to to generate the image and Muse is a two two State models so the first stage generates the 256 by 256 image and then in the second stage it's also um very similar to the first stage but it has two contexts it takes the text the generated low resolution generated tokens from the first stage and also the text embeddings and they does cross attention over both of them and then generates a high resolution image and one of the reasons we don't directly generate 512 512 images because sometimes the model can generate high frequency details so it's so instead we choose the two stage approach and the results with Muse model are really great these are some of the text to image samples this is some of the zero shot mask in painting and this is mask free editing like you just provide an image and just mask a bunch of tokens from that image and then just give any prompt and the model has the zero shot ability to edit it according to the prompt without using a mask yeah and yeah so it's it's all great but it's it's not as with like other Google text image models it's not open source which is just kind of not a good good news because it's when the models are not open source like the community can try out the models it's very hard to evaluate the models it's very hard to understand the capabilities of the model and for example evaluating models is notoriously hard like the current metrics like FID or clip score don't completely capture the capabilities of model or instead like if we give the model out to the community to try it out and like do different kinds of generation use different types of parameters and what notes we can like very easily assess the qualities of the model find out different biases one of the great examples is stable diffusion I mean like everyone has seen like explosion of stable diffusion and most of the the things that we can do the community has done with stable diffusion also most of these now go there a current like a lot of methods like tune a video or control net are now based on stable diffusion so it's because like that single model was open source a lot of new methods and applications came out of it so and so in the spirit of open reproduction at hugging Pace we have kick-started a project to openly reproduce Muse and we are collaborating with Robin Romberg the creator of stable diffusion and the open source Community to openly reproduce and train views model and then along the way share all artifacts code and learnings with the community so the community can build on top of it we have open source the code at this GitHub URL and all the intermediate checkpoints will be available on the Hub as as we train them so far we have two models just class conditional models to to test the approach and so far it's been going well and and the plan for the next couple months is to scale up the model and train it on bigger data sets so stay tuned and if you have any any feedback just go to the repo and share your suggestions or if you want to interact or if you want to join the efforts just feel feel free to just come in there it's it's going to be all open source so yeah that's it for me hello I I really liked the idea um it's it's very often that when you are like um when you are training masked models in text uh like you only can learn about the distribution of things but like you cannot really do the stuff like I don't know how it would work for um textbook like in painting and out painting with mask uh masking is like very intuitive and I really liked it I had this question but it's it was like a bit of a stupid question because like apparently when you do masking it's like it's it's it's like um masking in the like when you are masking patches and other stuff it's not like you don't have bi-directionality but you have non-directionality actually I think it's a I mean I'm sorry I didn't get the question so like you directly feed the so basically you you mask the tokens and then you directly feed the the thing itself like the masked image itself right yes yeah that was why I was saying it's not really bi-directional it's like non-directional actually like you don't really try to predict the next thing or like uh yeah it's a bit like birds yeah exactly so um we have like a couple of questions it's um Alex is asking once a token is unmasked can it be changed again say mask token Arrow token X Arrow of token y I don't know I think uh ideally it can but the way we train these models it's just like once a token is predicted there's we don't have a surety that the model will like predict another token for that position it's also one of the question I want to like do a bit of research on as part of open music like for example during training instead of just replacing a token with mask can we replace it with some random token so the during inference we can we predict the on like mask tokens again and then we predict those and another question is is it possible to use other sampling techniques than temperature based analogous to sampling form from language models type K top P or locally typical sampling yeah yeah it's definitely possible to use any of the top K or top B so this is the sampling is very similar to the way we do it in NLP cool for now we are just using temperature basically that's what was proposed in the paper and we wanted to kind of um train the models and compare with compare it with the paper so we kept is is it like did you try to like infer or are you still training is it hard to get right the the search parameters and like the temperature top carry and more it's like the default parameters sometimes work but to get best Generations I think you just need to kind of do a hyper parameter search like just choose a bunch of values for temperature top gate of p and then just evaluate the models cool um another question is does the future steps can impact already unmasked tokens or was it trained in a way that only unmasked tokens Can impact future tokens and not the opposite path that's an interesting question um in a way like there's no future step here because the model is bi-directional yeah it does attention or or all the tokens so both mask tokens can influence unmaster contains and unmasked tokens influence Mass tokens yeah and that's what I guessed as well so Christopher is asking can sometimes explain to me like I'm five half tokenization works for an image um that would be a different token itself but I'll try to see if I can find some slides so the tokenization I think we use uh Vector contacts vae or vector quantize again here it's basically like a vae but there's a discretization bottleneck um so the basic idea is like once you you have an encoder and then you get some sort of embeddings there's an embedding table and then the discrete tokens are the tokens which are kind of represent those embeddings and the model is Trend in such a way that uh um like as the model compresses the image we it's able to find the nearest embeddings using the tokens I'm not sure if that makes sense I don't have the slides for like any diagrams for a week you can now but I'm happy to expel it on the repo if you ask the question there um Andreas is asking how different can the unmasked king during inference be from unmasking during training how different [Music] um like I think yeah no I I just wanted to explain the question yeah I didn't get the question completely uh it's I think Andreas clarified it a little bit how different can the unmasking strategy during infants be from the unmasking strategy during the training how does this impact quality of the generated image I haven't really experimented with it and there are no details in the papers as well so it's um I guess we will have to find out cool and this was the last question thank you so much I really liked your presentation see you and I can have Sabrina now hello Sabina hey how's it going good good how are you doing all right um just came from a wild kayaking trip that had to be awarded thanks to weather conditions so we got to ride on a fast fast boat with Park Rangers and law enforcement delightful um but yeah so that's why I'm in this Airbnb with just my laptop so now funky you choose green magic thanks a lot for making slime to present today yeah for sure would you like to share a screen all right let me do that [Music] there we go do you see the slides yep perfect wonderful I will leave you with your presentation then good luck alrighty thanks hey so I'm Sabrina Mi okay I'm a grad student at Hopkins um hopefully finishing soon um I've kind of spent my time writing my thesis lately which is why I haven't been doing that much new stuff with Jax um and if you look at this a little tagline here uh it says the slides are almost two years old at this point um and the blog post that it's all based on is even older um but we're going to get to that in a second so the talk title is kind of a mouthful from stateful code to purify Jacks how to build your own Nelnet framework um and I guess it's somewhat of a different talk um than the ones that you've seen before uh in that I'm not really going to be describing interesting results New models um state of the art uh bleeding ads things um I really want to give you more of a an insight and a bit of like a uh hand holding to your own tinkering with some of the Frameworks that you're using in this workshop and other stuff you might see in the Jax ecosystem so um when Jax first really became a thing a topic people were interested in um it was often described as like numpy plus something so you might say it's numpy plus jit so an implementation that used to be on CPU like in this paper uh some students and I wrote back in the day um you could have like an SK learn implementation that was great 10 seconds on CPU but you're like well we're in the age of gpus can't we put this on the GPU and with Jax well that's super easy you just implement it call jit on it and run it on your accelerator and you get some nice speed ups that way in fact when I started messing around with it um I had this little thing where I was like okay how fast can I make lstms go because they weren't like entirely dead at that point just yet um and Jeff Dean actually retweeted me and I'm so happy this is the highlight of my career um so you can have a lot of fun just by making things fast with Jax okay fine um but then interesting for us is also of course the Grady I'm taking right if we want to do something with Radiance which we do that's kind of all we do in a way um it makes sense of Jax to think of Jax AS numpy Plus gradient taking um and this is not just your boring old gradients you can do uh pretty wild things like super easily um I don't know how many of you are familiar with reversible jump mcmc um if you're doing like Bayesian sampling you might have gone into mcmc and reversible jump is kind of a cool way to change dimensionalities but anyway I don't I don't really want to talk about that and you don't really need to look at all of this code um what I just want to highlight is that the math says you need the Jacobian determinant of a diffiomorphism and okay you look that up you implement it a little diffiomorphism that's just a function like this um okay it's literally just a python function how do I get the Jacobian determinant of that well with Jax I literally just call the Jacobian uh get this reversible Jacobian then I get the determinant and voila that's all I need pretty wild um like super convenient imagine having to do the math nobody wants to do that right um so that's cool um and then yeah you add your jid on top of the grad and you v-map it or P method and speed amazing um you can do some very fun things with that but what I want to talk about um here is more like if you're coming from the previous age of Frameworks so think Pi torch and eager mode tensorflow um once you look under the hood of how these things work it gets kind of weird you're like how where's the dot backward and house is actually done and so how do we actually Implement a neural NET Framework or even just train a neural net um with this sort of Technology to show the difference here's kind of what you would imagine it would look like in say apply torch if I take two tensors let's say w and x and I multiply them to create a new tensor y um what the sensor Y is actually going to have is not just the storage for the data not just the number 546. but it's also keeping track of where it came from it's keeping track that it was the result of this particular memory object that is the tensor called W in this particular memory object that's the tensor called X oops there goes the chat oh well um so this is kind of nice because now every tensor knows where it's coming from and so if you want to do a backward pass well every tensor just goes to his predecessors and so on right this is what we can see here once we are calling dot backward starting from the usually the last note but just the final note of your graph um we reverse mode Auto diff by following these connections and that's propagate the gradients at the top all the way down to the leaves where usually we're interested in the gradients of the weights so that's how we work in pytorch you build up your graph just by writing out the math you call Dot backward and whoops everyone everything has gradients wonderful um in Jax it's a little different because everything is pure and we're thinking a lot about transformations what do I mean by that um first of all what we're writing are pure functions like this I've called this function f very creative I know um and it does the same thing we've done before we just multiply these two tensors w and X together um we can draw a little picture okay it's like two things go in no one comes out wonderful um and now if I want to get the gradients in there what I do is I transform this function I take just a function f no values call Jack's grad on it and it gives me a new function I've called dftw um it still takes the same inputs so it takes the W and the X but now it gives me the gradient of what used to be the singular return value of x of f with respect to W and X or by default just with respect to W the first parameter of a function so I've drawn it this way okay so great now that's that that looks like what I remember from high school in math you do derivatives that's yeah you get a function now what um how do I how do I work with this right um well I evaluate this function using my current W and my current X I mean imagine W being weight and X being the data and I get the actual gradient value which is in this case 42. the gradient with respect to W and that's the same as it was here in the pi torch case that also gave me everything else okay I'll do it I transform the function I evaluate the transform function I get a number that's my gradient or I get a chance or that's my gradient great that's that's some nice toy examples um but but how are we actually going to go from these three liners to training a big neural network where it's not just a function f that multiplies its two inputs right or or said differently how can weights and gradients be stateful if our tensors can be mutated for those who are terminal the internet brain or the form it appears in their life um so as I said this talk is based on this blog post that I wrote in March yeah that's right March 2020 um it's like multiple lifetimes ago wasn't it um and so you might be like wow that's that's that's great that's really exciting um we're getting like the reheated leftovers um human ass is still relevant um and the answer is no it's it it never was relevant in the sense that it never really dealt with like oh state-of-the-art things this is idiomatic code this is cool results um it's really like our post you could have written a long time ago um just just about the theory of it um but something that people are now interested in because all these big Frameworks are out there and we want to maybe contribute to them or just understand what they're doing um and I found it really hard to understand what they were doing when I started out with Jax so that's kind of where this is coming from um hope after this talk um it'll all make a lot more sense to you and you're not gonna have to trial and arrow your way through it like I did so what are we going to accomplish again there will be no no shiny funny things to look at no eye candy there's no going to be any idiomatic allergens I will not claim that my coding style is the right one just kind of works and small enough that they shouldn't be too much wrong with it and we're also not going to do many of these cool things that I was talking about earlier we're like yeah you can batch distribute and compile um because it's kind of orthogonal to what we're trying to achieve and you could always add it on top while we are going to do though and finally I'm saying something positive we are going to be looking at the dead simple example of linear regression good old MX plus b um we're going to do it statefully so like the pytorch way try and train it on some data and you know okay just reminder how we used to think the software then we're going to be trying to move this example to Jax while dealing with this notion of Purity that is required for a gradient taking there I'm gonna have a little aside you can get a coffee if you're not that interested um where we're kind of like um what if actually we could do this cool thing that people don't do but you could do um and then we're going to be working towards the actual foundations of a framework that kind of does the work that we did in step two for us and finally we're going to find out oh that's that's what these Frameworks do all along big surprise sorry I spoiled it so here's what a mod should look like in pytorch I don't know how many people still remember I feel like I have forgotten um I literally had to like watch a recording of this talk yesterday last night being like Oh yeah yeah that's the stuff that's right um but what you see is essentially that modules own their weights there are these parameters which are just wrapped tensor objects um and they're assigned to this unit this lstm cell and then in the forward method you use um when you build this execution graph this tape that we've seen before um so far so good um once you do like your actual training Loop what you're going to be doing is create your instance instantiate all the weights and then call Dot forward which uses the current weight dot backward which starts this whole process of going backwards to the graph right and writing out all the gradients everywhere and then once we have gradients on all our Leaf notes apply them to the data like if this was the middle SGD um and that's basically how you would do it um unlike the most Bare Bones way with play torch if you mutate the gradients by calling last step backward and you mutate the weights by calling parameter data minus equals and both of these things are no longer allowed once we move to Jax now I'm going to switch to the notebook part and I hope you can you can still see this and I hope my runtime hasn't died oh no but this runtime has died well that's no good everything seems to be fine except for the runtime okay wonderful wonderful all right by the way I I wanted to just comment that I wish they actually would teach deep learning Frameworks like this in schools like I've been enjoying your talk a lot um okay um thank you I try to get it to comic comic code ligatures again this is one of my favorite coding fonts but I guess it only lasts up one cell depressing what can what are we gonna do some some people are asking if you could make the phones bigger oh right let me do that yeah how's that thank you yeah it looks good thank you all right so this is the the thing that I just said um we're doing linear regression we're doing it on Torch we're being a very boring um we have our weights they're being owned blah blah blah production users self dot something we have a loss function in this case just root mean squared error that again uses um the parameters that are owned by the module blah blah blah number comes out let's create this object um and then let's test it out for a second does it predict for the number three well it predicts zero which makes sense we've initialize the things to zero right zero times zero plus zero zero times three plus zero is zero wonderful how how exciting um Let's do let's do some training though um we're going to create a very small data set as literally just one example that three should be map to seven [Music] um and since we already know that three is mapped to zero what do we expect the root mean squared error to be that should be seven um so far so good um so now that we have a loss that works we can do a little bit of training so let's make the data set a little bigger so we're not like totally under defined um you can make out in your head what M and V fit this one um but it should be fairly straightforward let's just train this we get the loss we backward we go through the parameters we have the gradient blah blah blah we do this little step that when I learned this stuff I always love to forget and wonder what my networks didn't train properly and then hopefully we're going to see this prediction come closer to the true answer which is seven and let's see voila yes wonderful it goes there it doesn't really hit seven or step size is too large we should be smart with it but we don't want to be smart we just want to be smart enough for this talk that's going to be enough okay so that's the pie charts example just to refresh your memory you can come back from your coffee break now because we're going to be looking at it in Jax this is how you might try to implement this thing in checks you're like okay I'm gonna make a class I'm gonna make weights and stuff they're going to be owned by the module predict and then we do the rmse yup yup cool very similar nothing interesting happening here um let's let's see what happens uh on this one we should predict zero yup and we should have a loss of seven yep so far so good it's a very direct translation um but now we're going to start running into issues because now we need to get gradients right what do we take the gradient off well we usually take the gradient of the loss function right the loss function is RMS e um and things get a little weird here because this function is not at all Pure or at least um it doesn't make sense to think of it as a pure function because it takes the entire object and differentiating with respect to a python object seems a little weird so before I go even further into this discussion of Purity that just sounds very strange and we're not on the same page what it means let's go look at some examples um I've written down three functions here and already given away the answers the first one is pure second one impure and the third one pure again and I want to illustrate some things on these functions basically a pure function the idea is something goes in something goes out that's all the function does it doesn't mess with any of the other stuff that's around in the ambient atmosphere so for the first function that certainly is true we put a number in and twice the number comes out nothing has been mutated or illegally accessed mutation is why the second function is no longer pure because if I call it it is going to actually change the list that I'm giving as an argument so to illustrate this let me Define a list down here L um and call impure function 1 on l well and let me do it again results has changed in fact it has changed uh oops not excess l in fact we can see very clearly that L is no longer the same so this function is like evil it took something and then it didn't do stuff with it and got something else out no it took the list that I gave it and messed with it that's a no no um it's also kind of easily spotable by the fact that you call this function twice with what was supposed to be the same input and you've got different results a lot of sketchy stuff on the other hand um this function over here this pure function also gets a list and also does some list magic but it only appends internal state so it defines this new variable Ys and that variable you can append stuff that's fine you can return the variable for all you care you know as long as you don't mess with your arguments and don't mess with anything that's outside of the function everything is fine and we can see that that's true um like calling the function multiple times gives me the same result and the original list is still the same it was that that looks pretty pure to me so that's nice um let's do some more examples uh here's some more impure functions that are a little more interesting to look at and while we look at these functions I'm going to look at the chat since my phone fell down okay cool um yeah what's wrong with these functions well for one I've already talked about this ambient State nonsense and this impure function too very obviously mutates Global stage right does this it drags the number executions variable from the global State into the function changes it and then sure it'll join something but it sounds something evil inside we can see that that's true just by executing it a bunch of times we're like yeah I mean technically it Returns the same number every time but it changes this state it changes This Global variable that's bad that's not okay so now here's a question for you what about this third function impure function three uh sorry that's not that's not my example just yet um input function three I can I can quickly show you that's also sketchy um because call it sure it gives me something but then uh if I call it with a different Global state or executions I said right if I call it with a different Global state it suddenly gives us a different answer so that's bad and we don't want that um I was going to say something what was I going to say who knows kind of been that important okay so Global state that don't mess with it don't access it okay it's fun I'll be getting to my sorry on one laptop screen I have like my copy paste pad really tiny to the side so I can't see what's coming up this is very exciting um okay here we go final examples thank God right back to some fun Jack stuff in a second um here's another Classic this impure function does a lot of bad things that are somewhat if you've ever used like Haskell or some other functional language like you know all this stuff and you know these things are a little tricky at first like you know printing something out it changes the global State what I see in my console window or whatever is part of the global state so this is already changing things asking for the user input um that's definitely impure like I could I could put in something different every time and the results of this function might change that's no good the function must only take in information that it was explicitly given in its arguments and finally this one here we're asked for the current system time that's obviously also a bad idea right because that's going to change from execution to execution we're again accessing something we're not allowed to access so now one last function mpo function 5. which constraint does this violate what's wrong with it um it's impure I've given that much away um but what is it actually doing that's kind of a yeah let's see I'm gonna I'm gonna wait a few seconds uh to see if someone wants to DARE guess in the chat um I also don't know if I see the chat it won't go down yes okay yeah yeah that's the that's the that's right Wow first try damn I like when people struggle you know I like when they get things wrong it's much more fun um but yes so it acts as Global stage accessing the the seed and the state of the RNG and it also mutates the state of the random number generator right it advances to the next stage so it can get a new Fresh random number um so it kind of violates both these things of taking in things it wasn't supposed to look at and taking and having effects that it wasn't supposed to happen like the the the nature enthusiasts among us might be familiar with the take nothing but memories leave nothing but Footprints and that's kind of kind of applicable here in fact um one mnemonic that I found super useful um for understanding Purity and for having like a very quick and easy test is this one the function you're implementing could easily be a lookup table or it could easily be cached then it's pure this is obviously not true for this one right if I if I give it like a nice cash decorator uh wow now my function is cached and all the randomness is gone um so that's not true same for these things right I can't just cache user input or Cache the printing because it has to happen again when I call this function on the other hand or do we have a nice pure function over here this function I could totally cash that right it's never going to change doesn't really depend on anything else so that's a really nice and easy way to tell whether your function is pure if you can convert it to a lookup table or you can cache it it's pure okay so how are we going to do this is the time jump back to the slides yet who knows we'll find out um how are we going to build the last function that is pure we do it like this we are still using this weird object oriented skeleton that we had up here right where it owned its own way it's um and it's used Eminence method but just like our Pure example we're past the list and then created our own list that only mutated things inside this will be a pure loss function right yes it creates a stateful class and then it mutates the stateful class and then it calls the stateful class that accesses its own thing statefully but in the bigger scheme of things if I just look at this function well it just takes W and V and axis wise and put some in there and then calls this result and then the linear regressor is destroyed afterwards no trash is left we don't look at anything we're not supposed to look at so this loss function would be pure and again the gives us gives us the result that we wanted you can use this last function in fact to already start training this thing again let's expand our data a little to make it less underdefined um here's a loss function we pass the stuff in it's very beautifully pure um and we take the gradients um with with respect to the first two parameters that's what this is for because usually grad only gives you the gradient for the first parameter but we want gradients for w and for B so we tell it give us gradients for both of these things and that's what's coming out DW DB okay and now uh we have this this class my regressor that we kind of keep around statefully that's okay we're kind of we're not letting Jack smell Jax doesn't deal with this this is like our own record keeping um and we can mutate it again doing this very simple SGD um and then predict with it and we see again trans beautifully gets close to seven that's all we want that's all we were asking for okay now if you give me a second I'm sneakily looking at my phone to see when I'm jumping back to so okay no not for a while neat you get to see more notebooking love it um okay so this is nice and all but um it's kind of silly to have this this my regressor class and not use its function and just use it to hold the W and the B for our own bookkeeping right that's what we're doing here like we might as well not bother you know screw this you could wear this and let's store our parameters in a nice dictionary you know there's no much of a difference except the dictionary feels lighter and better and more modern and Sleek love that so um with this we might change our method slightly to now have our Pure loss function not get W and B or in other words not enumerate parameter by parameter of our parameter but only have this one dictionary params and then still use it as before get the W entry get the B entry calculate the loss function easy peasy um well function it's a little different now because we still have access and wires to the big data set that's fine cool so we've saved ourselves a bit of typing right up here that's that's nice I'd love to save myself some typing um but it gets even better because now if you're asking yourself okay well this is a pure function but this is not a number how what the hell is a gradient for a dictionary so what do you think if I do this if I call jax.grad on the loss function now again we're only interested in the first parameters because we don't have to Consolidated it that gives me the gradient function and I evaluate it on my current parameters and my data set what's the type of gradients going to be now again we'll open up the chat I'll give you a second yes exactly it's a dictionary as well dictionary that uses the same keys that we've been giving W and B and the values are now the gradients calculated by jacks and these are the same gradients we have before and we're going to see them a bunch more throughout this notebook so um get familiar with these numbers but yeah um it just works like that this is amazing right like in my loss function I don't have to enumerate all the parameters in here I don't have to enumerate all the parameters in here I don't have to enumerate all the parameters the only place I ever really deal with the parameters is when I use them in my loss function that's pretty cool um and pretty convenient at that so now with this we can we can again do this thing where we update our current parameters go through all these things and if you're looking at this you might already be like why don't you just Loop through all the keys and and not write things out line by line and patience we will get there in just a second what matters is loss goes down 7.5 goes to 6.1 amazing we've done it the machine has learned so I said uh we're going to get there in a second and the second has come um if we get our gradients like this and we know their dictionary now right um we could Define such a function called like an update combiner this is really if you think about it like a very Bare Bones Optimizer skeleton um which would take a single parameter a single gradient like these could be still be tensors they don't have to be scalars but a single item um give it a learning rate let's be fancy so we really see what's happening and it just Returns the new value of the parameter which is parameter minus that gradient just as before and now what we can do is we can map this update combiner over the tree in this case it's just a dictionary but it we're going to see in a second why it could be called a tree um by passing it to this function called jax.treymap um together with the params and the gradients and it is going to do this work for us of going through all the keys in this dictionary or all the leaves in this tree um updating them and giving us a new object what is the type of my parents kind of obvious at this point or at least we would hope that it would be something that we like um it is a dictionary that contains a new parameter estimates and indeed these new parameter estimates that come out of the update again have made our laws go down wow this is great to update a question and chat that I just saw um by trees I'm going to mention in a second uh what does weak type mean and I wondered about the same thing when I checked out this notebook this morning I was like all right let's change some Jackson's last time this wasn't here um that's got to do with the type promotion semantics so when I when I put in something like um yeah over here I say w is 0.0 what is that type something floating point right but is it like 16 32 64-bit maybe it's even a complex number um just by writing 0.0 that's not really specified and we could say like oh this is actually an ah an array with um b-type something like oops can't type and so on but um if we're not going to do that um then Jack just like I don't really know about this one let's say this is a weak type for now that's that's a gist of it maybe someone with more uh experience working with this can weigh in in the chat too um but yeah I wanted the very same thing um this is something else I don't think so right cool so let's move on with that we've made the loss go down we've used jax.tree map which used to be called jax.3 multimap you'll still see that in the old blog post and stuff but it's it's a small change um you really just use it the same way so that's wonderful um now that we've done some individual last steps let's let the machine do this for us we do a few steps of STD um again this is all the same stuff we've seen we take the gradient of the loss function this we don't actually have to use every Loop we could just have done it once and then cache the function but who cares go slow the world is beautiful or something and then do the treatment update and this is a little awkward that we still have this predict method coupled to our linear regressor and now we have to manually transfer the waves from the dictionary into this object that holds our magic methods but we're going to get there in a second what matters is think trains great so we so with this all done it is time for an aside uh so if you want to catch a break this is it this is going to be a fun little fun little thing that might help you understand pie trees a little better and someone already mentioned this in the YouTube comments and I think it's one of the most amazing and beautiful things in Jax um you're wondering maybe why could we do this thing with dictionaries how does Jax know that it has to go through dictionaries like some were in Jax it says numbers are numbers tensors are numbers dictionaries are close enough to numbers they'll have like numbers in them not good enough um and the whole API is talking about Pi trees as things that have numbers in them that we can take out and do math with so you might really think of it as a tree um this makes most sense in when you think of like maybe a nested dictionary writer dictionary that contains some dictionaries that contain some dictionaries that contain some tensors and so Jax is somehow able to maintain the structure do things on the tensors in between like take radians and then recreate this tree with that structure with the new numbers inside and the Beautiful Thing is we can do that too um it kind of lets us do this where we can register new types of Pi tree nodes numbers tensors lists dictionaries and our class linear regressor we're going to make that a pi treature the way you do that is you call Jack's free util register proprietary node um give it the class that you want to register and then you give the two functions these functions are pure because all beautiful things are pure um are flattened and unflat the flattened one basically takes this complex object this Pi tree node that we have and converts into a representation that sort of tees us apart the two parts of this structure whatever the whatever the structure may be one part we call the leaves is really a tuple or a list of all the tensors in there the stuff with the numbers the stuff that Jax is going to be dealing with and the other part concerns the greater structure like if there's some structural information that is not numerical that we can't differentiate through or anything like that but that is important in this in this object and must be kept for operations like tree map in this case since we only really have these two scalos here I mean we just make Leaves the Tuple and there's no auxiliary information really um when we return that Tuple that's what you do for flatten and then for unflattened you basically do the same thing in Reverse you go from leaves and auxiliary information back to the object However unfortunately due to historical baggage the argument order is reversed right here so be careful with that um again since we don't really use aux I've just underscored it um we use the leaves let's say it's Tuple wmd we know this because we've created it ourselves and then we create a new linear regressor assign the right values and return it and so with these functions Jax can now when it is asked to differentiate through a linear regressor uh it can be like all right I don't know what the that is but uh let me call flatten on this thing and then I get a tuple with like numbery things that I do know what to do with and auxiliary things that I'm just not gonna touch all right you do your thing on the numbers now you call unflattened put it all together and voila you have a new linear regressor this is quite beautiful magic in my opinion um I greatly enjoy registering Pine notes Apache nodes all the time and with this we can now do cool things um our last function can now take the regressor as an object and do the rmse um and okay that's fine we were able to do that before but we weren't able to do before is get gradients so if I now call jax.grad last function well we get a lot we get a function that gives us the gradient with respect to the first parameter which is my regressor my class so what's the type of gradients what do you think yeah it leaves the pi tree structure basically untouched in the arcs thing all that matters is that leaves is a tuple or at least um I don't know if that matters but I'll just make it a tuple um actually anyway a tuple our list and um that's all that's all you need to do a tree yes it's a pie tree um but what kind of pine tree is it Ah that's right it's a linear regressor right here access created this class for us using its magic and really using our hard work of writing the flattened and unflattened methods that's pretty cool I think I don't think so that's fine you can stay um and we can look at these uh properties of this object these these members of our class and yes we get these familiar gradients all has worked out at before we can even do tree map with this thing because tree map again works with pi trees so we just combine the original regressor with the gradient regressor and we get a new regressor and under that new regressor the loss has gone down so that's that's pretty much my side um if you thought this was cool and you want to see how far to take this um you probably want to find a way to automatically write flattening on setting Logic for people's new classes so they don't have to do this maybe you use type annotations maybe you hijack the methods kind of like pie torch oops um and you want to do something with initializers some Randomness and if you follow this through you might get this uh this little framework that Sasha rash and I sketched in 2020 where we were basically following this idea and seeing how far we could get um it's not really like been a thing at all it's really just a fun project and we kind of stalled on making batch normal work because you are statefulness um but if you want to take a look at it you know that's there oh it could be fun all right so um now let me think what do I switch back okay cool I think it's time to switch back to the slides oh less typing less room for mistakes love to see it um apple pie tree I love that cool so um what we're going to be doing next is forget about all these poetry stuff that's what I'm saying it was inside it was fun we had a good time we had laughs but you know we still wake up the next morning sober up and now it's time to get back to work um and actually think about what we can do to make this dictionary based writing or whatever we're doing just a little bit nicer what we want to do is make the user's life easy the user's gotten very used to writing these stateful methods like in pi torch um I'm gonna say yeah fine write your staple methods I'll clean up your mess after how do we clean up the mess here's what we're going to be doing and this is really the key idea behind these Frameworks and behind this talk um we're going to have this class regressor as before it owns its weights um and you know if this was ply torch then this function f that takes in an X yeah I would use the weights never access State outside of the function it should not be pure and return it's why wonderful what we're going to be doing is let the user write a class like that but actually do a little trick what we want to be doing is um called a purifying regressor and I totally made this up but I thought it sounded pretty cool um where we don't really start out with a weight that's kind of weird there's none here we'll deal with that later um and the user writes this function f just as before saying oh I access my weight blah blah blah things will be fine um and then we are somehow tasked with the job of creating a purified version of that app that now actually takes in the parameters alongside the data first thing it does is assigns of the W to this information it just took in and so now what happens when we next call this dirty function f is yeah we pass in the X we get out of Y but when this function f acts as a self.w it doesn't actually access anything from outside the function because we've just defined solve the W to be the argument itself so the information flow doesn't leave the function we don't actually violate purity because we've sort of tricked this function it's running in a kind of sandbox if you will so after that I mean we might not bother but I like to keep things clean um we just set self to W Back To None and we return the Y that came out of the dirty function and as you can see all the arrows are inside the function nobody writes anything anywhere they're not allowed to and nobody reads anything from whether we're not allowed to so that's pretty cool um and this is kind of how I want it to be um I want to have this sort of fancy magic artisanal base module class the same way we have in all our favorite Frameworks that's going to have to do this work for us this purification somehow I don't see it in my code that's that's that's annoying and I want to Define my impure function like in this case we have an impure rmsa that takes in the data so far so good um and then uses its parameters now this is not quite itself.w uh which maybe you could make it work but it'll be much easier to compromise on saying fine we have like a method called get param that sort of acts like self.w right we're also going to pass this method and initializer that looks kind of weird like we're mixing things that wasn't supposed to be mixed um but we're going to see why that's going to be really useful later on and so now these W and B that we get from get param we are legally allowed to assume are actually just normal tensors nothing weird about them so we use them in our math and return on value and something's going to happen here there's got to be some magic that's going to come from the fancy magic artational base module class um and maybe it's called purify or something we'll see we'll figure it out you know we want to have a function that basically purifies this impure function into a purified function and I'm already going to spoil we're going to get an initializer out of this that's why we mixed it up in there so we don't have to enumerate like okay here's our dictionary W zero V is zero um let's just let's just like this function take care of it you know why why make why work when the framework can work for you and then once we get my params that's just going to be a dictionary just as before so you see that get param is actually accessing this dictionary um spoilers and so we can call a purified thing with it we can do gradients of the purified thing and call that with it we can do Jack Stream app and get new parameters just as we've seen before okay so let's implement this thing first off to the notebook it says and that's what we will do here is my draft of what I want this to look like we have the purifying regressor it uses self.params as opposed to self.getparam because you know let's build this up simple let's not do the initializer but just yet keep it simple um okay this is great this is like what I want to write because I'm boring and I never want to learn new things and I want to write pytorch code and supposedly that's a good it's a good thing to do we have params initially set to none and then this purified version of the function set self.params to the params that were passed called the emperor function that tries to use self.params but jokes on you there's nothing actually you're getting that isn't coming into this function to begin with and then we zero it out in the end just to be clean and return the result initialize parameters by hand for now get the value and graph because value and grad is nice for what we want to be doing when we look at the loss and the gradients and the purify function you will see is actually working because yeah we can get gradients through it and they are the same gradients we've been seeing all along in this notebook that's delightful um let's see how can we make this this basically now do the fancy or Crystal base magic framework I would like to not buy too much code well this one's going to be a little bigger so I'm gonna decrease the font size let me know what's a good uh compromise like maybe this will work so um since we're going to be doing um pure execution and initializing I was like super uber hyper careful and anal about it and just added two flags to this class that says are we running in pure mode are we running an initialization launch ideally they should not be true at the same time thank you there's better ways to do this but yeah whatever it's good enough um parameters are known as before here's our function get param that I hinted at in the slides and that is now going to depend on which mode we're in if we are initializing then we're going to check do we already have that parameter in our uh in our params dictionary if so nothing we need to do if not um create it with the initializer um that's why we pass in the initializer as the third argument to get param okay and then return the parameters so the function can continue if we're running pure well then we basically assume that the self.parents has already been hijacked and we can just use it so sure return self.parents you know good for you enjoy um and if we're not running and initializing or in pure mode well then we're impure and we don't like being impure so we raise an exception wow so much guardrail so much safety feels very German um now we're going to have a bunch of functions um this one is basically the wrapper for the initialization logic it gives us a function called initialize params that is returned that sets us to initializing mode gives parameters an empty dictionary calls the method that it was passed then when the method that Mayberry will have been impure and we assume it to be impure is done we said initializing to false we return self.params and we set self the params To None to not risk leaking information the bureau method works much the same way and we've already seen this before now where we create a new pure method given an impure method called method we said I'll be pure to true we set the parameters up and we call the method that is impure about jokes on you we've already hijacked subtle params and again zero it out after and set the mode back to the original one the final one is the API that we want to expose to the user so maybe I should have underscored the other ones but I hope you'll forgive me um given a method I want us to return a tuple that consists of the purified version of the method and an initialized params version of the method that basically is going to initialize all the parameters and give us the parameter dictionary based on all the parameters that were used in this method that's a lot of stuff like this is a big sell and if you're a little overwhelmed I I understand but we're pretty much done at this point with with the hard lifting now we're just going to use it um this is pretty much like what I wanted to see in the slide we have an impure RSM rmse use self.get param we give it an initializer that's very simple and boring we have an impure project that also uses wmp and gets get param the same way um and then we call this purify that we just defined up there and it gives us the purified one and the initialize one we call my params and let me actually add this here because why not we can see what's happening and that's going to be oh yeah yeah oh I didn't execute the cell that would have been useful so we should see the initial parameters pop up in a second that should be zero zero um yep except we didn't have to write it down the framework data for us um and then we could get the loss before the update we get the gradients we get the laws after the update it's all beautiful as it was before except we had even less code to write and that's kind of the beauty of it all um again obligatory uh mechanization of turning this into like a loop our our loss goes down our correctness goes up you know and that's really all I'm going to show you from this notebook to end this talk with what kind of recap from what we built we built this thing up here right where we have a get param and we have this purify and then we call initialize and purified rmse um but when we now look up over the landscape of all these beautiful Frameworks blossoming uh in the distance we've kind of re-implemented exactly what they're doing and this is Steve Minds Haiku which was my framework uh framework of choice at the time it literally has a version a method called get parameter and it gives it the initializer and then it calls our purify called transform and it gives you a transformed in it and transformed apply which is yeah which is our our initialize params and our purify function exactly the same thing a little bit more Belton whistles to deal with Randomness and state and all that stuff but you know you get the idea um yeah or um more more appropriate maybe now these days uh where where Flex has become a big fan if you look at Flagstaff linen um it kind of uses this annotator and out compact um to do the transformation for us so the class already uh turns its call method into the init and apply pair but apart from that again the same thing happens we have this method called shelf.param gets this name that accesses the dictionary that's Underneath It All We have an initializer in this case we have a shape too and then we have the internet and applying same as we did well again with Randomness handling which we skipped over but you know that's that's that's that's really all there is to it so um that's really all I have to say about this we've kind of started From First principles and tried to figure out how could one do this and mess around with some ideas and in the end we came up with something that's pretty much exactly what the Frameworks are doing because you know we all come up with the same idea so often enough so if you're interested in more details about this more complications a little bit more background maybe a slower explanation um I'll just more code this is all on my blog post and there's a notebook um version of it too and you can see the slides of this talk and the recording from back when and can compare how I'm doing today you can follow me on Twitter if you want I'm not really active on Twitter anymore these days it's one of the addictions that I've actually managed to replace so yeah hey any questions and I'll be here yeah there is one question because you answered them throughout this Stream So if we separate the parents dictionary from the architecture how do we know when the parameters are broken by a change and architecture mm-hmm so you would kind of Wonder like how do I tell whether a dictionary is compatible with the architecture that it's using mind I mean I really think that there's no no good answer to that um like you could do something like okay let's take the um the architecture and like hash it and store the hash in the dictionary that that'll be a quick test um maybe that's what some Frameworks actually do I don't actually know I feel like I vaguely remember something like that but I'm not sure but even if not I mean this is one of those things that will break it'll just break if you access a parameter that isn't there it breaks um I I am of the opinion that things breaking is a good thing I'd rather things break than me um but that's kind of what I would say this is all about it gets a little weirder when you have like parameters that are in there but aren't used but then the optimizer uses them so then things get wacky um but yeah I guess yeah like hashing the architecture somehow is probably the best way to go and again I don't know what exactly Frameworks are doing and Alex is asking with the gradient being another linear regressor is calculating the second derivative as simply as just calling loss function on the regressor again um hmm I don't see why not um so what did we first Define this [Music] is all torch never mind uh you should go a little beloved I think [Music] um yeah here we go cool um and this is what we have what we have so let's add another cell um let's just call it Jack so grab twice it's the worst that can happen right um and this is a function I'm gonna call it separately with I will pressure x's and y's gonna happen well we'll find out the functions all right so what we need is uh okay um so it's not really a gradient in this point right it would be um like a Jacobian or something wait no no Second Assassin something like that um yeah yeah I wonder you know I don't actually know no idea maybe it's something obvious and easy and I just don't see it right now but that's a fun one I don't think yeah we have any more questions uh thanks a lot for this presentation I loved it I will definitely check out the follow-up one uh Alex is saying you need to pass the returns gradients into same function returned by Jack scrub the loss function um I don't know how it is that should be the same thing um like because loss function gives you a scalar jax.grad gives you uh or executing jax.grad gives you the linear regressor so the type is messed up anyway yeah you can't get a gradient with respect to a regressor you could do it with respect to all the individual leave node entries or something um but I think I think really in this case since we are dealing with like more than scalars we are dealing with like multiple Dimensions the second gradients would not be gradients um they would and should be something something in quadratic so Christian saying grad loss function grad loss function um what would that mean why would you run it over like that that will be something but it would not be the second derivative yeah yeah this is different than the second derivative I think well I'm not the Jax Pro so and he is a good one it's acting on function without argument so that doesn't work Okay Christian is saying not saying this is right and he says but it will give something it will give something that's for sure yeah I think the types will work out on that one that's true um that looks like a good workout yeah anyway um there is a Hessian function by the way he says that makes sense I mean I don't know how it works in this case because we have multiple parameters that are not in the same tensor so you get some weird cross tensors like the object wouldn't really match right because you have these interactions between W and B and they're stored in different tensors so you need like some WP and BW tensors so I think it would get a little it will get a little messy but I might be completely wrong with that who knows so thanks a lot for coming today I really enjoyed your talk and um so um thanks a lot I'm going to take you off now see ya um we are going to have another round of talks on 17th and that will be the day three of our talks and we have speakers from Google brain we will have a Boris who created the dalimini in the previous Jax sprints and stay tuned and also set your reminders for the day three of them thank you so much for uh chiming in and I will see you around thank you